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

结果:

相关推荐
Python×CATIA工业智造1 小时前
Frida RPC高级应用:动态模拟执行Android so文件实战指南
开发语言·python·pycharm
onceco1 小时前
领域LLM九讲——第5讲 为什么选择OpenManus而不是QwenAgent(附LLM免费api邀请码)
人工智能·python·深度学习·语言模型·自然语言处理·自动化
狐凄2 小时前
Python实例题:基于 Python 的简单聊天机器人
开发语言·python
悦悦子a啊3 小时前
Python之--基本知识
开发语言·前端·python
笑稀了的野生俊5 小时前
在服务器中下载 HuggingFace 模型:终极指南
linux·服务器·python·bash·gpu算力
Naiva5 小时前
【小技巧】Python+PyCharm IDE 配置解释器出错,环境配置不完整或不兼容。(小智AI、MCP、聚合数据、实时新闻查询、NBA赛事查询)
ide·python·pycharm
路来了5 小时前
Python小工具之PDF合并
开发语言·windows·python
蓝婷儿5 小时前
Python 机器学习核心入门与实战进阶 Day 3 - 决策树 & 随机森林模型实战
人工智能·python·机器学习
AntBlack6 小时前
拖了五个月 ,不当韭菜体验版算是正式发布了
前端·后端·python
.30-06Springfield6 小时前
决策树(Decision tree)算法详解(ID3、C4.5、CART)
人工智能·python·算法·决策树·机器学习