как получить значения прогнозов для тестовых изображений с помощью YOLO

Я новичок в моделях обнаружения объектов, я обучил модель YOLO с размером изображения 600, а затем нанес удар по получению файла значений тестового прогнозирования best.pt, после чего у меня есть файл . Вот мой код:

      from ultralytics import YOLO
model=YOLO("runs/detect/train16/weights/best.pt")

Вот мои тестовые данные:

      img_fName   img_w   img_h
0   07e0f900-e1d2-4e3c-a812-651a156d7463.jpeg   4032    2268
1   98171e52-0def-4de8-b9ad-23cc3a1f1a2d.jpeg   2448    3264
2   30da361c-282e-4a8c-a714-4de4d5407696.jpeg   2988    5312
3   780ce034-878f-41fc-a23f-5432987ce64a.jpeg   3000    4000
4   6e7e3e96-62a4-46e3-aaa2-ce88f2949f6a.jpeg   4128    3096

Здесь я пытаюсь получить значения тестового прогноза:

      for index, row in test.iterrows():
image =  row['img_fName']
img_w = row['img_w']
img_h  = row['img_h']
# Load the image
print(path)
image = Image.open(f'test_images_phase1')

# Resize the image to the expected size
image = F.resize(image, (img_h, img_w))

# Convert the image to tensor
image_tensor = F.to_tensor(image)

# Add a batch dimension
image_tensor = image_tensor.unsqueeze(0)

# Run inference
with torch.no_grad():
    predictions = model(image_tensor)

# Extract bounding box coordinates, class labels, and confidence scores
boxes = predictions[0]["boxes"].tolist()
labels = predictions[0]["labels"].tolist()
scores = predictions[0]["scores"].tolist()

# Print the prediction values for the image
for box, label, score in zip(boxes, labels, scores):
    x, y, x_max, y_max = box
    class_label = model.classes[label]

    print(f"Image: {image}")
    print(f"Bounding Box: {x}, {y}, {x_max}, {y_max}")
    print(f"Class: {class_label}")
    print(f"Confidence: {score}")
    print()

0 ответов

Другие вопросы по тегам