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

一、前言:

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

安装包:

pip install super_gradients
pip install omegaconf
pip install hydra-core
pip install boto3
pip install stringcase
pip install typing-extensions
pip install rapidfuzz
pip install Cython
pip install pycocotools
pip install onnx-simplifier

二、步骤:

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

三、代码:

from PIL import Image

import torch
from super_gradients.training import models

device = 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 = 1
for 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 Image

import torch
from super_gradients.training import models

device = 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")

结果:

相关推荐
pzx_00110 分钟前
【内积】内积计算公式及物理意义
数据结构·python·opencv·算法·线性回归
一丝晨光11 分钟前
逻辑运算符
java·c++·python·kotlin·c#·c·逻辑运算符
ForRunner12314 分钟前
使用 Python 高分解决 reCAPTCHA v3 的指南
数据库·python·microsoft
从懒虫到爬虫1 小时前
YOLOv5/v8 + 双目相机测距
yolo·qq767172261·双目测距
躺平的花卷1 小时前
Python爬虫案例六:抓取某个地区某月份天气数据并保存到mysql数据库中
数据库·爬虫·python·mysql
虚拟搬运工1 小时前
Python类及元类的创建流程
开发语言·chrome·python
学步_技术2 小时前
Python编码系列—Python原型模式:深克隆与高效复制的艺术
开发语言·python·原型模式
FL16238631292 小时前
[数据集][目标检测]不同颜色的安全帽检测数据集VOC+YOLO格式7574张5类别
人工智能·yolo·目标检测
阿利同学2 小时前
热成像目标检测数据集
人工智能·目标检测·机器学习·目标跟踪·数据集·热成像目标检测数据集·机器学习 深度学习
Desire.9842 小时前
Python 数学建模——灰色关联度分析
python·数学建模·灰色关联度