Python はターゲット検出に YOLO_NAS_S モデルを使用し、予測された被写体画像を保存します

I.はじめに:

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

2. 手順:

  1. 必要なライブラリとフレームワークをインストールします。OpenCV、PyTorch、torchvision がインストールされていることを確認してください
  2. YOLO_NAS_S モデルのウェイト ファイルをダウンロードし、モデルをロードします
  3. 画像の前処理を実行します。入力画像ごとに、モデルが受け入れられる形式に変換し、正規化する必要があります。
  4. モデルをターゲット検出に使用し、予測結果を取得する
  5. 予測結果を分析し、予測されたメインピクチャを保存します

3. コード:

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

結果:

おすすめ

転載: blog.csdn.net/xun527/article/details/132537183