一、前言：

使用 YOLO_NAS_S 模型进行目标检测，并保存预测到的主体图片

安装包：

pip install super_gradientspip install omegaconfpip install hydra-corepip install boto3pip install stringcasepip install typing-extensionspip install rapidfuzzpip install Cythonpip install pycocotoolspip install onnx-simplifier

二、步骤：

  1. 安装所需的库和框架。确保已经安装了 OpenCV、PyTorch 和 torchvision
  2. 下载 YOLO_NAS_S 模型的权重文件，并加载模型
  3. 进行图像预处理。对于每张输入图像，需要将其转换为模型可接受的格式，并进行归一化处理
  4. 使用模型进行目标检测，并获取预测结果
  5. 解析预测结果，并保存预测到的主体图片

三、代码：

from PIL import Imageimport torchfrom super_gradients.training import modelsdevice = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")model = models.get("yolo_nas_s", pretrained_weights="coco").to(device)out = model.predict(r"D:\Desktop\tp.png", conf=0.6)predictions = out[0]# 提取预测框对应的主体图像并保存num = 1for bbox in predictions.prediction.bboxes_xyxy:x1, y1, x2, y2 = bbox[:4]# 每个预测框的坐标image = Image.open(r"D:\Desktop\tp.png")cropped_image = image.crop((x1, y1, x2, y2))# 根据坐标裁剪图像output_path = f"output_{num}.jpg"cropped_image.save(output_path)# 保存裁剪后的图像num += 1

被检测的图片：

预测主体效果：

如果在原图的基础上查看代码如下：

from PIL import Imageimport torchfrom super_gradients.training import modelsdevice = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")model = models.get("yolo_nas_s", pretrained_weights="coco").to(device)out = model.predict(r"D:\Desktop\tp.png", conf=0.6)out.save("save_folder_path")

结果：