使用 Python 和 YOLO-NAS 检测图像和视频中的对象

创建虚拟环境

我们需要做的第一件事是创建并激活虚拟环境，这可以通过以下命令完成：

python3 -m venv env
source env/bin/activate

完毕！接下来我们需要创建一个需求文件并安装依赖项。

安装依赖项

创建 requirements.txt 文件并添加以下内容：

imutils
torchinfo
torch
super-gradients

要安装依赖项，请运行以下命令：

pip install -r requirements.txt

现在我们终于可以开始编码了。

检测图像中的对象

首先，我将向您展示如何检测图像中的对象，打开一个名为 image.py 的文件并添加以下导入：

import gc
import os
import torch
from torchinfo import summary
from super_gradients.training import models

import argparse

接下来我们需要创建一个方法来检测图像中的对象，代码如下：

def detect_objects(image):
    device = torch.device("cpu")
    model = models.get("yolo_nas_s", pretrained_weights="coco").to(device)
    out = model.predict(image, conf=0.6)

    out.save("predictions");

    del model
    gc.collect()
    torch.cuda.empty_cache()

上述方法采用图像补丁、检测对象，然后将输出保存到预测目录。最后我们清理一下。

最后我们需要一个main方法：

if __name__ == "__main__":
    os.makedirs("predictions", exist_ok=True)

    ap = argparse.ArgumentParser()
    ap.add_argument("-i", "--image", required = True, help = "Path to input file")
    args = vars(ap.parse_args())

    detect_objects(args["image"])

main 方法接受一个图像参数，然后将其传递给前面定义的 detect_objects 方法。

最后，该命令可以通过以下方式运行：

python image.py -i [path to image file]

完成后，您应该会看到类似于以下内容的内容：

检测视频中的对象

检测视频中的对象与上面的代码完全相同，除了我们传递视频文件而不是图像文件。

由于我已经解释了上面代码的作用，所以我将在此处提供完整代码：

import gcimport osimport torchfrom torchinfo import summaryfrom super_gradients.training import modelsimport argparse
def detect_objects(video):    device = torch.device("cpu")    model = models.get("yolo_nas_s", pretrained_weights="coco").to(device)    out = model.predict(video)          
    out.save("predictions") 
    del model    gc.collect()    torch.cuda.empty_cache()
if __name__ == "__main__":    os.makedirs("predictions", exist_ok=True)
    ap = argparse.ArgumentParser()    ap.add_argument("-v", "--video", required = True, help = "path to input video")    args = vars(ap.parse_args())

    detect_objects(args["video"])

请注意，我将图像参数更改为视频参数。

这可以通过以下方式运行：

python video.py -v [path to video file]

处理后的视频现在应该出现在您的当前目录中。

如果你想使用GPU而不是CPU，把cpu换成cuda