yolo_tracking+设置检测区域+过线计数

最新推荐文章于 2024-08-08 08:30:32 发布

莫言过

最新推荐文章于 2024-08-08 08:30:32 发布

阅读量1.6k

点赞数 27

文章标签： YOLO 目标跟踪 pytorch ubuntu

本文链接：https://blog.csdn.net/qq_43979221/article/details/137708887

版权

最近使用yolov8结合追踪算法实现过线计数功能，查找了很多资料，没找到什么免费的教程，实现过程耗费了不少时间，记录一下实现过程，分享一下。
实验环境：Ubuntu 22.04.3 LTS、Driver Version: 535.104.05、CUDA12.1、cudnn8.9.6、conda24、torch2.2.1、python3.10

car (2)

0 前言

实验代码基于mikel-brostrom的开源项目yolo_tracking进行更改，要求python>=3.8，torch>=2.2.1，太低版本的Ubuntu安装不了新驱动，旧驱动安装不了新CUDA，旧CUDA安装不了torch2.2.1，故需要先对环境进行一系列升级。满足环境要求可以直接跳过此节，原本实验室两台机器Ubuntu版本是18.04和20.04，下面记录一下环境升级过程。

0.1 升级Ubuntu系统

18.04需要先升级到20.04，成功后在执行一遍下方的代码升级到22.04：

sudo apt-get update
sudo apt-get upgrade
sudo apt-get install update-manager-core
sudo reboot
sudo do-release-upgrade -d

在升级包的时候遇到如下报错：
E: Failed to fetch https://deb.termius.com/dists/stable/InRelease 403 Forbidden E: The repository ‘https://deb.termius.com stable InRelease’ is not signed. N: Updating from such a repository can’t be done securely, and is therefore disabled by default. N: See apt-secure(8) manpage for repository creation and user configuration details.
尝试了一系列方法没有解决，后面直接把这个包删了：

sudo apt-get --purge remove termius
sudo apt autoremove

再次运行upgrade显示有包没有更新：
0 upgraded, 0 newly installed, 0 to remove and 1 not upgraded.
后来尝试换源、使用aptitude、sudo apt install --only-upgrade code等一系列方法没有解决，
陆续报错：
The following packages have been kept back: code
Unable to correct problems, you have held broken packages.
最后又把这个包直接删了，把有问题的包删除后就可以升级了。

apt list --upgradable
sudo apt-get remove code
sudo apt autoremove

参考链接：
https://blog.csdn.net/leisurelen/article/details/106014823

0.2 升级驱动

nvidia-smi查看驱动版本
在这里插入图片描述
上方分别为当前驱动版本以及所支持的最高CUDA版本，查看可升级的驱动版本并更新至推荐的版本：

sudo ubuntu-drivers devices
sudo apt-get install nvidia-driver-535

在这里插入图片描述
（不建议使用sudo ubuntu-drivers autoinstall自动安装，没有尝试，见其他帖子说的）

参考链接：
https://zhuanlan.zhihu.com/p/643773939

0.3 安装CUDA

如果之前安装了CUDA需要先卸载：

#如果安装过则/usr/local/ 文件夹下包含cuda文件夹
cd /usr/local/cuda/bin
sudo ./cuda-uninstaller
sudo rm -rf /usr/local/cuda-11.2  #11.2替换为自己的cuda版本号

卸载后去官网选择合适的版本，复制对应的runfile命令下载：
在这里插入图片描述

sudo wget https://developer.download.nvidia.com/compute/cuda/12.1.1/local_installers/cuda_12.1.1_530.30.02_linux.run
sudo sh cuda_12.1.1_530.30.02_linux.run

上方命令行第一行下载对应的CUDA版本，第二行安装
安装时会提示Existing package manager installation of the driver found. It is strongly recommended that you remove this before continuing，选择Countinue，回车后输入accept，然后选择具体要安装什么，因为之前已经安装过驱动了，这时候要把驱动前面的X去掉（移动到对应行敲回车），然后移动到Install回车开始安装。
在这里插入图片描述

安装好之后要设置环境变量：

sudo vim ~/.bashrc  # 配置系统环境：打开系统环境配置文件，按i编辑，然后输入：
# 将cuda的头文件和库文件路径加入系统环境，因为有软链接，可以直接使用 /usr/local/cuda
# CUDA
export PATH=/usr/local/cuda/bin:$PATH
export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH
export CUDA_HOME=$CUDA_HOME:/usr/local/cuda
# 按Esc退出编辑，输入：wq！保存退出

source ~/.bashrc #更新环境变量
nvcc -V  #测试命令

在这里插入图片描述
我一开始使用vim ~/.zshrc设置环境变量，但是发现每次打开新的窗口都需要使用source ~/.zshrc更新环境变量，后来使用 vim ~/.bashrc就没有这个问题了。

参考链接：
https://blog.csdn.net/aiboom/article/details/123307212?spm=1001.2014.3001.5506
https://feihu.me/blog/2014/env-problem-when-ssh-executing-command-on-remote/

0.4 安装cudnn

在cudnn官网注册账号并下载合适的版本：
在这里插入图片描述

下载的文件类型为archive.tar.xz （可以本地下载后把文件复制到服务器），进入下载的路径，解压后把对应的文件复制到对应的CUDA文件夹：

tar -xvf cudnn-linux-x86_64-8.9.6.50_cuda12-archive.tar.xz

sudo cp cudnn-linux-x86_64-8.9.6.50_cuda12-archive/include/cudnn*.h /usr/local/cuda/include
sudo cp -P cudnn-linux-x86_64-8.9.6.50_cuda12-archive/lib/libcudnn* /usr/local/cuda/lib64
sudo chmod a+r /usr/local/cuda/include/cudnn*.h /usr/local/cuda/lib64/libcudnn*

查看cudnn版本：

cat /usr/local/cuda/include/cudnn_version.h | grep CUDNN_MAJOR -A 2

在这里插入图片描述

参考链接：
https://blog.csdn.net/hymn1993/article/details/132240184?spm=1001.2014.3001.5506

1 配置环境

1.1 创建虚拟环境

使用conda创建虚拟环境：

conda create -n torch221_py310 python==3.10
conda activate torch221_py310

conda remove -n torch221_py310 --all  #卸载虚拟环境命令

在创建时，conda提示我版本太老，推荐我使用conda update -n base -c defaults conda升级conda版本，但是我运行后没能成功升级，报错The environment is inconsistent, please check the package plan carefully The following packages are causing the inconsistency
尝试换源、回退版本、增加参数–repodata-fn=repodata.json等方法无效，最后手动删除了这些包就可以正常更新了（参考链接）。

创建好环境后去官网复制pytorch安装命令：
在这里插入图片描述

conda install pytorch==2.2.1 torchvision==0.17.1 torchaudio==2.2.1 pytorch-cuda=12.1 -c pytorch -c nvidia

1.2 下载代码并安装依赖

首先去github下载代码：

git clone https://github.com/mikel-brostrom/yolo_tracking.git

下载后原作者推荐使用poetry安装（个人推荐使用第二种方法）：

cd yolo_tracking
pip install poetry
poetry install --with yolo  # installed boxmot + yolo dependencies
poetry shell  # activates the newly created environment with the installed dependencies

我使用poetry安装的时候一直卡着不动，查阅了各种资料也没有解决，最后使用了下面的安装方法（所需依赖见pyproject.toml）：

cd yolo_tracking
pip install -e .
git clone https://github.com/mikel-brostrom/ultralytics.git    #对应的ultralytics版本为8.0.228
cd ultralytics
pip install -e .
cd ../

至此，yolo_tracking环境已经配置好了，代码已经可以正常运行了（运行时代码会自动在外网下载模型，如果不能科学上网需要提前手动把yolo模型和reid模型下载下来放到yolo_tracking文件夹）。

python tracking/track.py --source 'car.mp4' --yolo-model 'yolov8n.pt' --reid-model 'osnet_x0_25_msmt17.pt' --save --tracking-method bytetrack

在这里插入图片描述

在这里插入图片描述
测试结果如下：

检测视频上传到了视频栏目，原视频为car，原始模型测试结果为car(1)，更改过线计数后的测试结果为car(2)，yolo模型没有经过训练，效果一般。后续yolo模型经过训练可以达到较好的性能，在代码中没有发现在哪里用到了reid模型，但是运行追踪代码需要指定reid模型（代码中默认为osnet_x0_25_msmt17.pt），实测只训练yolo模型就可以达到不错的追踪效果。

参考链接：
https://blog.csdn.net/qq_37950540/article/details/88901496
https://github.com/mikel-brostrom/yolo_tracking
https://docs.ultralytics.com/tasks/detect/
https://kaiyangzhou.github.io/deep-person-reid/MODEL_ZOO
https://www.bilibili.com/video/BV1te4y1Q7mN/

2 yolo_tracking函数调用过程

运行跟踪代码：

python tracking/track.py --source 'car.mp4' --yolo-model 'yolov8n.pt' --save --tracking-method bytetrack

track.py重载了on_predict_start函数，用来加载自己的追踪器；run函数首先创建了一个YOLO类对象（继承Model类），然后调用Model类的track函数（ultralytics/ultralytics/engine/model.py），然后调用add_callback传入boxmot的跟踪器：

    yolo = YOLO(
        args.yolo_model if 'yolov8' in str(args.yolo_model) else 'yolov8n.pt',
    )
    results = yolo.track(args)
    yolo.add_callback('on_predict_start', partial(on_predict_start, persist=True))
    yolo.predictor.custom_args = args

track函数首先判断self.predictor（ultralytics/engine/predictor.py）对象是否有trackers属性，如果没有则调用register_tracker（ultralytics/trackers/track.py）创建跟踪器，然后在return时调用self.predict函数

def track(self, source=None, stream=False, persist=False, **kwargs):
    """
    Perform object tracking on the input source using the registered trackers.
    Returns:
        (List[ultralytics.engine.results.Results]): The tracking results.
    """
    if not hasattr(self.predictor, 'trackers'):
        from ultralytics.trackers import register_tracker
        register_tracker(self, persist)
    kwargs['conf'] = kwargs.get('conf') or 0.1  # ByteTrack-based method needs low confidence predictions as input
    kwargs['mode'] = 'track'
    return self.predict(source=source, stream=stream, **kwargs)

predict函数设置预测器并创建模型，如果是cli模式则在return调用predict_cli，如果不是的后半句没看懂，不过估计功能差不多

def predict(self, source=None, stream=False, predictor=None, **kwargs):
    """
    Perform prediction using the YOLO model.
    Returns:
        (List[ultralytics.engine.results.Results]): The prediction results.
    """
    if source is None:
        source = ASSETS
        LOGGER.warning(f"WARNING ⚠️ 'source' is missing. Using 'source={source}'.")

    is_cli = (sys.argv[0].endswith('yolo') or sys.argv[0].endswith('ultralytics')) and any(
        x in sys.argv for x in ('predict', 'track', 'mode=predict', 'mode=track'))

    custom = {'conf': 0.25, 'save': is_cli}  # method defaults
    args = {**self.overrides, **custom, **kwargs, 'mode': 'predict'}  # highest priority args on the right
    prompts = args.pop('prompts', None)  # for SAM-type models

    if not self.predictor:
        self.predictor = (predictor or self._smart_load('predictor'))(overrides=args, _callbacks=self.callbacks)
        self.predictor.setup_model(model=self.model, verbose=is_cli)
    else:  # only update args if predictor is already setup
        self.predictor.args = get_cfg(self.predictor.args, args)
        if 'project' in args or 'name' in args:
            self.predictor.save_dir = get_save_dir(self.predictor.args)
    if prompts and hasattr(self.predictor, 'set_prompts'):  # for SAM-type models
        self.predictor.set_prompts(prompts)
    return self.predictor.predict_cli(source=source) if is_cli else self.predictor(source=source, stream=stream)

predict_cli主要为调用self.stream_inference，进入正式预测流程

def predict_cli(self, source=None, model=None):
    """
    Method used for CLI prediction.

    It uses always generator as outputs as not required by CLI mode.
    """
    gen = self.stream_inference(source, model)
    for _ in gen:  # running CLI inference without accumulating any outputs (do not modify)
        pass

stream_inference函数为预测的流程代码：
首先调用setup_model设置模型，然后调用setup_source加载数据，模型Warmup后使用run_callbacks函数调用on_predict_start函数（如果使用boxmot的追踪模型则调用yolo_tracking/tracking/track.py，使用yolov8的追踪模型则调用/ultralytics/trackers/track.py）构建追踪器，然后对每个batch进行处理：
profilers[0]：调用preprocess进行数据预处理，
profilers[1]：调用inference进行预测，
profilers[2]：调用postprocess进行后处理，返回results列表，然后使用run_callbacks调用on_predict_postprocess_end（/ultralytics/trackers/track.py）进行追踪；
然后在 for i in range(n) 下对每张图片进行处理，判断是否需要save或者show并调用对应函数（write_results等）进行处理；
最后打印相关日志信息。

def stream_inference(self, source=None, model=None, *args, **kwargs):
    """Streams real-time inference on camera feed and saves results to file."""
    if self.args.verbose:
        LOGGER.info('')

    # Setup model
    if not self.model:
        self.setup_model(model)

    with self._lock:  # for thread-safe inference
        # Setup source every time predict is called
        self.setup_source(source if source is not None else self.args.source)

        # Check if save_dir/ label file exists
        if self.args.save or self.args.save_txt:
            (self.save_dir / 'labels' if self.args.save_txt else self.save_dir).mkdir(parents=True, exist_ok=True)

        # Warmup model
        if not self.done_warmup:
            self.model.warmup(imgsz=(1 if self.model.pt or self.model.triton else self.dataset.bs, 3, *self.imgsz))
            self.done_warmup = True

        self.seen, self.windows, self.batch, profilers = 0, [], None, (ops.Profile(), ops.Profile(), ops.Profile())
        self.run_callbacks('on_predict_start')

        for batch in self.dataset:
            self.run_callbacks('on_predict_batch_start')
            self.batch = batch
            path, im0s, vid_cap, s = batch

            # Preprocess
            with profilers[0]:
                im = self.preprocess(im0s)

            # Inference
            with profilers[1]:
                preds = self.inference(im, *args, **kwargs)

            # Postprocess
            with profilers[2]:
                if isinstance(self.model, AutoBackend):
                    self.results = self.postprocess(preds, im, im0s)
                else:
                    self.results = self.model.postprocess(path, preds, im, im0s)

            self.run_callbacks('on_predict_postprocess_end')
            # Visualize, save, write results
            n = len(im0s)
            for i in range(n):
                self.seen += 1
                self.results[i].speed = {
                    'preprocess': profilers[0].dt * 1E3 / n,
                    'inference': profilers[1].dt * 1E3 / n,
                    'postprocess': profilers[2].dt * 1E3 / n}
                p, im0 = path[i], None if self.source_type.tensor else im0s[i].copy()
                p = Path(p)

                if self.args.verbose or self.args.save or self.args.save_txt or self.args.show:
                    s += self.write_results(i, self.results, (p, im, im0))
                if self.args.save or self.args.save_txt:
                    self.results[i].save_dir = self.save_dir.__str__()
                if self.args.show and self.plotted_img is not None:
                    self.show(p)
                if self.args.save and self.plotted_img is not None:
                    self.save_preds(vid_cap, i, str(self.save_dir / p.name))

            self.run_callbacks('on_predict_batch_end')
            yield from self.results

            # Print time (inference-only)
            if self.args.verbose:
                LOGGER.info(f'{s}{profilers[1].dt * 1E3:.1f}ms')

    # Release assets
    if isinstance(self.vid_writer[-1], cv2.VideoWriter):
        self.vid_writer[-1].release()  # release final video writer

    # Print results
    if self.args.verbose and self.seen:
        t = tuple(x.t / self.seen * 1E3 for x in profilers)  # speeds per image
        LOGGER.info(f'Speed: %.1fms preprocess, %.1fms inference, %.1fms postprocess per image at shape '
                    f'{(1, 3, *im.shape[2:])}' % t)
    if self.args.save or self.args.save_txt or self.args.save_crop:
        nl = len(list(self.save_dir.glob('labels/*.txt')))  # number of labels
        s = f"\n{nl} label{'s' * (nl > 1)} saved to {self.save_dir / 'labels'}" if self.args.save_txt else ''
        LOGGER.info(f"Results saved to {colorstr('bold', self.save_dir)}{s}")

    self.run_callbacks('on_predict_end')

profilers[0]：preprocess函数，进行数据预处理（之后设置检测区域需要更改此函数），转换数据格式，进行归一化等；

def preprocess(self, im):
    """
    Prepares input image before inference.

    Args:
        im (torch.Tensor | List(np.ndarray)): BCHW for tensor, [(HWC) x B] for list.
    """
    not_tensor = not isinstance(im, torch.Tensor)
    if not_tensor:
        im = np.stack(self.pre_transform(im))
        im = im[..., ::-1].transpose((0, 3, 1, 2))  # BGR to RGB, BHWC to BCHW, (n, 3, h, w)
        im = np.ascontiguousarray(im)  # contiguous
        im = torch.from_numpy(im)

    im = im.to(self.device)
    im = im.half() if self.model.fp16 else im.float()  # uint8 to fp16/32
    if not_tensor:
        im /= 255  # 0 - 255 to 0.0 - 1.0
    return im

profilers[1]：inference函数，调用self.model进行推理预测和追踪；

def inference(self, im, *args, **kwargs):
    """Runs inference on a given image using the specified model and arguments."""
    visualize = increment_path(self.save_dir / Path(self.batch[0][0]).stem,
                               mkdir=True) if self.args.visualize and (not self.source_type.tensor) else False
    return self.model(im, augment=self.args.augment, visualize=visualize)

profilers[2]：postprocess函数，预测后处理（nms操作），创建并返回results列表；

def postprocess(self, preds, img, orig_imgs):
    """Post-processes predictions and returns a list of Results objects."""
    preds = ops.non_max_suppression(preds,
                                    self.args.conf,
                                    self.args.iou,
                                    agnostic=self.args.agnostic_nms,
                                    max_det=self.args.max_det,
                                    classes=self.args.classes)

    if not isinstance(orig_imgs, list):  # input images are a torch.Tensor, not a list
        orig_imgs = ops.convert_torch2numpy_batch(orig_imgs)

    results = []
    for i, pred in enumerate(preds):
        orig_img = orig_imgs[i]
        pred[:, :4] = ops.scale_boxes(img.shape[2:], pred[:, :4], orig_img.shape)
        img_path = self.batch[0][i]
        results.append(Results(orig_img, path=img_path, names=self.model.names, boxes=pred))
    return results

write_results函数进行结果绘制，调用result.plot(**plot_args)在图片绘制bbox，调用result.save_txt保存txt信息

def write_results(self, idx, results, batch):
    """Write inference results to a file or directory."""
    p, im, _ = batch
    log_string = ''
    if len(im.shape) == 3:
        im = im[None]  # expand for batch dim
    if self.source_type.webcam or self.source_type.from_img or self.source_type.tensor:  # batch_size >= 1
        log_string += f'{idx}: '
        frame = self.dataset.count
    else:
        frame = getattr(self.dataset, 'frame', 0)
    self.data_path = p
    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
    result = results[idx]
    log_string += result.verbose()

    if self.args.save or self.args.show:  # Add bbox to image
        plot_args = {
            'line_width': self.args.line_width,
            'boxes': self.args.show_boxes,
            'conf': self.args.show_conf,
            'labels': self.args.show_labels}
        if not self.args.retina_masks:
            plot_args['im_gpu'] = im[idx]
        self.plotted_img = result.plot(**plot_args)
    # Write
    if self.args.save_txt:
        result.save_txt(f'{self.txt_path}.txt', save_conf=self.args.save_conf)
    if self.args.save_crop:
        result.save_crop(save_dir=self.save_dir / 'crops',
                         file_name=self.data_path.stem + ('' if self.dataset.mode == 'image' else f'_{frame}'))

    return log_string

plot函数（ultralytics/engine/results.py）绘制结果：
首先对传入的图片创建Annotator对象（ultralytics/utils/plotting.py），通过调用Annotator的box_label函数绘制bbox和label。

def plot():#复制了主要代码
    annotator = Annotator(
        deepcopy(self.orig_img if img is None else img),
        line_width,
        font_size,
        font,
        pil or (pred_probs is not None and show_probs),  # Classify tasks default to pil=True
        example=names)
    # Plot Detect results
    if pred_boxes and show_boxes:
        for d in reversed(pred_boxes):
            c, conf, id = int(d.cls), float(d.conf) if conf else None, None if d.id is None else int(d.id.item())
            name = ('' if id is None else f'id:{id} ') + names[c]
            label = (f'{name} {conf:.2f}' if conf else name) if labels else None
            annotator.box_label(d.xyxy.squeeze(), label, color=colors(c, True))
    return annotator.result()

on_predict_start创建跟踪器（如果使用yolo的追踪算法则调用/ultralytics/trackers/track.py中的on_predict_start函数）：

def on_predict_start(predictor, persist=False):
    """
    Initialize trackers for object tracking during prediction.

    Args:
        predictor (object): The predictor object to initialize trackers for.
        persist (bool, optional): Whether to persist the trackers if they already exist. Defaults to False.
    """

    assert predictor.custom_args.tracking_method in TRACKERS, \
        f"'{predictor.custom_args.tracking_method}' is not supported. Supported ones are {TRACKERS}"

    tracking_config = TRACKER_CONFIGS / (predictor.custom_args.tracking_method + '.yaml')
    trackers = []
    for i in range(predictor.dataset.bs):
        tracker = create_tracker(
            predictor.custom_args.tracking_method,
            tracking_config,
            predictor.custom_args.reid_model,
            predictor.device,
            predictor.custom_args.half,
            predictor.custom_args.per_class
        )
        # motion only modeles do not have
        if hasattr(tracker, 'model'):
            tracker.model.warmup()
        trackers.append(tracker)

    predictor.trackers = trackers

on_predict_postprocess_end函数调用predictor.trackers[i].update(det, im0s[i])进行追踪：

def on_predict_postprocess_end(predictor: object, persist: bool = False) -> None:
    """
    Postprocess detected boxes and update with object tracking.

    Args:
        predictor (object): The predictor object containing the predictions.
        persist (bool, optional): Whether to persist the trackers if they already exist. Defaults to False.
    """
    bs = predictor.dataset.bs
    path, im0s = predictor.batch[:2]

    for i in range(bs):

        det = predictor.results[i].boxes.data.cpu().numpy()

        if len(det) == 0:
            continue
        tracks = predictor.trackers[i].update(det, im0s[i])
        if len(tracks) == 0:
            continue
        idx = tracks[:, -1].astype(int)
        predictor.results[i] = predictor.results[i][idx]
        predictor.results[i].update(boxes=torch.as_tensor(tracks[:, :-1]))

至此，整个算法的流程基本就结束了，之前不太熟悉调试，花费了很多时间，并且还有好多地方没弄明白，后来打断点调试了一遍看了一遍代码的运行流程就比较清晰了，建议学习一下调试方法，能节省很多精力。

3 设置检测区域

更改BasePredictor类（ultralytics/engine/predictor.py）的preprocess函数：

def preprocess(self, im):
    """
    Prepares input image before inference.

    Args:
        im (torch.Tensor | List(np.ndarray)): BCHW for tensor, [(HWC) x B] for list.
    """
    # mask for certain region，设置检测区域并画框
    left,upper,right,lower=0,165,640,265#设置检测区域，四个值分别为左上角和右下角像素坐标
    #1,3 分别对应左上，右下，归一化
    hl1 , wl1 = upper / im[0].shape[0] , left / im[0].shape[1] #监测区域高度距离图片顶部比例 左部比例
    hl3 , wl3 = lower / im[0].shape[0] , right / im[0].shape[1] # 监测区域高度距离图片顶部比例 左部比例
    
    masked_images = []
    for image in im:
        # 对每个图像进行遮罩覆盖操作
        height, width, _ = image.shape
        mask = np.zeros((height, width), dtype=np.uint8)
        mask_start_y, mask_start_x = int(height*hl1), int(width*wl1) # 遮罩区域起始位置的y坐标,x坐标
        mask[mask_start_y:int(height*hl3), mask_start_x:int(width*wl3)] = 255
        masked_image = cv2.bitwise_and(image, image, mask=mask)
        # 检测区域画框
        cv2.rectangle(image, (mask_start_x, mask_start_y), (int(width*wl3),  int(height*hl3)), (0, 255, 0), 2)
        masked_images.append(masked_image)
	# 主要是在原来的函数增加了上面的代码，原理就是对要检测的以外的区域置黑
	# 参考链接：https://blog.csdn.net/qq_39740357/article/details/125149010

    not_tensor = not isinstance(masked_images, torch.Tensor)
    if not_tensor:
        masked_images = np.stack(self.pre_transform(masked_images))
        masked_images = masked_images[..., ::-1].transpose((0, 3, 1, 2))  # BGR to RGB, BHWC to BCHW, (n, 3, h, w)
        masked_images = np.ascontiguousarray(masked_images)  # contiguous
        masked_images = torch.from_numpy(masked_images)

    masked_images = masked_images.to(self.device)
    masked_images = masked_images.half() if self.model.fp16 else masked_images.float()  # uint8 to fp16/32
    if not_tensor:
        masked_images /= 255  # 0 - 255 to 0.0 - 1.0
    return masked_images

参考链接：
https://blog.csdn.net/qq_39740357/article/details/125149010

4 增加参数，实现检测区域中间画线

增加参数hidecls设置保存结果时不显示种类：

1、修改tracking/track.py中的parse_opt函数，增加：
parser.add_argument('--hidecls', default=False, action='store_true', help='')

2、调用results = yolo.track时传入形参：
hidecls=args.hidecls,

2、在ultralytics/cfg/default.yaml中增加hidecls，默认不隐藏：
hidecls : False

3、在BasePredictor类的__init__函数中初始化hidecls：
self.hidecls = False

4、在write_results函数中的字典plot_args中增加键值对：
'hidecls' : self.args.hidecls,

5、在ultralytics/engine/results.py的plot函数形参列表中增加：
hidecls=False,

6、更改Plot Detect results的name：
name = ('' if id is None else f'{id} ') + (names[c] if not hidecls else '')

增加参数whereline记录线的位置，增加参数numlabel记录视频要显示的信息：

1、首先在BasePredictor类的__init__函数初始化：
self.whereline = None
self.numlabel = None

2、在上方的preprocess函数设置whereline 的值
left,upper,right,lower=0,165,640,265
# 设置whereline为检测区域中间，位置可以调整 
self.whereline = [(0, int(upper+(lower-upper)/2)), (int(im[0].shape[1]), int(upper+(lower-upper)/2))]

3、在write_results函数的字典plot_args中增加键值对：
		if self.args.save or self.args.show:  # Add bbox to image
            plot_args = {
                'line_width': self.args.line_width,
                'boxes': self.args.show_boxes,
                'conf': self.args.show_conf,
                'labels': self.args.show_labels,
                'hidecls' : self.args.hidecls,
                'whereline' : self.whereline,
                'numlabel' : self.numlabel,
                }
            if not self.args.retina_masks:
                plot_args['im_gpu'] = im[idx]
            self.plotted_img = result.plot(**plot_args)

4、参数plot_args被传入函数result.plot(**plot_args)，然后在plot函数形参列表中增加：
whereline=None,
numlabel=None,

增加box中心画点功能:

#在plot函数的# Plot Detect results代码块增加画点代码：
				drawboxes=d.xyxy.int().cpu().tolist()
                minX, minY, maxX, maxY = drawboxes[0][:4]
                midpoint = (int((minX + maxX) / 2), int((minY + maxY) / 2))
                annotator.drawpoint(midpoint, radius=5, color=(0, 0, 0), thickness=-1)

在polt函数的return 前增加画线代码：

# 画检测区域中间线
annotator.draw_line(whereline[0], whereline[1], (0, 255, 255), 4)
# 左上角显示数目
annotator.puttxt(numlabel,(50, 50),  2, (0, 0, 255), 5)
return annotator.result()

在Annotator类（ultralytics/utils/plotting.py）中增加画线、画点、信息输出函数：

def draw_line(self, x1, x2, color=(0, 255, 255),cuxi= 4):
    cv2.line(self.im, x1, x2, color=color,thickness=cuxi)

def drawpoint(self,box,radius=5, color=(0, 0, 0), thickness=-1):
    cv2.circle(self.im, box, radius, color, thickness)

def puttxt(self,text, p1, x=2,color=(0,0,255), y=5):
    cv2.putText(self.im, text, p1, cv2.FONT_HERSHEY_SIMPLEX, x, color, y)

5 过线计数

在BasePredictor类（ultralytics/engine/predictor.py）中增加计数用的函数：

def tlbr_midpoint(self,box):
    minX, minY, maxX, maxY = box
    midpoint = (int((minX + maxX) / 2), int((minY + maxY) / 2))  # minus y coordinates to get proper xy format
    return midpoint

def intersect(self,A, B, C, D):
    return (self.ccw(A, C, D) * self.ccw(B, C, D))<=0 and (self.ccw(A, B, C) * self.ccw(A, B, D))<=0

def ccw(self,A, B, C):
    if( (C[1] - A[1]) * (B[0] - A[0]) > (B[1] - A[1]) * (C[0] - A[0])):
        return 1
    elif((C[1] - A[1]) * (B[0] - A[0]) < (B[1] - A[1]) * (C[0] - A[0])):
        return -1
    return 0

def vector_angle(self,midpoint, previous_midpoint):
    x = midpoint[0] - previous_midpoint[0]
    y = midpoint[1] - previous_midpoint[1]
    return math.degrees(math.atan2(y, x))

在stream_inference函数的for batch in self.dataset:前增加初始参数：

from collections import Counter
from collections import deque
import math
#计数初始参数
paths = {}
angle = -1
up_count = 0
down_count = 0
for batch in self.dataset:

然后在for i in range(n):中增加计数代码：

for i in range(n):
    self.seen += 1
    self.results[i].speed = {
        'preprocess': profilers[0].dt * 1E3 / n,
        'inference': profilers[1].dt * 1E3 / n,
        'postprocess': profilers[2].dt * 1E3 / n}
    p, im0 = path[i], None if self.source_type.tensor else im0s[i].copy()
    p = Path(p)
	
	#numcount
    if self.results[i].boxes.id is not None:
        # print(self.results[i].boxes)
        for _,perbox in enumerate(self.results[i].boxes):
            # print("j ",_)
            perid=perbox.id.int().cpu().tolist()[0]
            percls=perbox.cls.int().cpu().tolist()[0]
            track_id = perid
            # print(perid,' ',percls,' ')
            minX, minY, maxX, maxY = perbox.xyxy.int().cpu().tolist()[0][:4]
            midpoint = (int((minX + maxX) / 2), int((minY + maxY) / 2))

            origin_midpoint = (midpoint[0], im0.shape[0] - midpoint[1])  # get midpoint respective to botton-left
            if track_id not in paths:
                paths[track_id] = deque(maxlen=2)
                paths[track_id].append(midpoint)
                paths[track_id].append(midpoint)
            if (self.ccw(midpoint,self.whereline[0], self.whereline[1])!=0):
                paths[track_id].append(midpoint)
            previous_midpoint = paths[track_id][0]
            origin_previous_midpoint = (previous_midpoint[0], im0.shape[0] - previous_midpoint[1])
			# 根据角度的正负判断是上行还是下行
            if self.intersect(paths[track_id][1], previous_midpoint, self.whereline[0], self.whereline[1]) :
                angle = self.vector_angle(origin_midpoint, origin_previous_midpoint)
                if angle > 0:
                    up_count += 1
                if angle < 0:
                    down_count += 1
            if len(paths) > 200:
                del paths[list(paths)[0]]
    # 赋值numlabel ，在视频左上角显示
    self.numlabel = "amount: {} ({} up, {} down)".format(str(down_count+up_count),str(up_count), str(down_count))
    
    if self.args.verbose or self.args.save or self.args.save_txt or self.args.show:
                        s += self.write_results(i, self.results, (p, im, im0))
    ...

需要注意本文实现的是计数所有过线的物体，上行和下行分开计数，总数为上下行数量之和；在检测时设置单一物种检测即可实现单一种类计数；若要实现各个种类分别计数需要把上方的几个参数设置为数组，然后修改代码，按照percls分别计数即可。

# 设置参数--classes 2，只检测类别号为2的物体（类别号同训练模型用的yaml文件中的类别号，yolov8训练方法和yolov5一样，可以参考下方帖子或者网上搜教程
python tracking/track.py --source 'car.mp4' --yolo-model 'yolov8n.pt' --tracking-method bytetrack --device 0 --save --classes 2

参考链接：
https://blog.csdn.net/qq_43979221/article/details/130360949
https://xugaoxiang.blog.csdn.net/article/details/118773692
https://blog.csdn.net/zengwubbb/article/details/113422048
https://blog.csdn.net/sxj731533730/article/details/119951580
https://github.com/zengwb-lx/yolov5-deepsort-pedestrian-counting/blob/master/person_count.py
https://blog.csdn.net/qq_51248362/article/details/134102209
https://blog.csdn.net/zhaocj/article/details/132800296

注：
本人根据此博客从配置环境重新完整执行了一遍流程，car(2)是最后得到的检测视频；有错误的地方欢迎大家指正交流共同进步。