Python uses the YOLO_NAS_S model for target detection and saves the predicted subject image

I. Introduction:

Use the YOLO_NAS_S model for target detection and save the predicted subject image

Installation package:

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

2. Steps:

1. Install required libraries and frameworks. Make sure you have installed OpenCV, PyTorch and torchvision
2. Download the weight file of the YOLO_NAS_S model and load the model
3. Perform image preprocessing. For each input image, it needs to be converted into a format acceptable to the model and normalized
4. Use the model for target detection and get prediction results
5. Analyze the prediction results and save the predicted main picture

3. Code:

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

Detected pictures:

Predict the subject effect:

 If you view the code on the basis of the original image as follows:

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

result: