Saving yolo Model Result as String

Question

I have a model, best.pt, that I'd like to run. It takes an image as input. It classifies an object in this image, a fruit.

from PIL import Image
from ultralytics import YOLO

# Load the pre-trained model
model = YOLO('best.pt')
# Load the input image
input_image = Image.open('fruit.jpeg')
# Pass the image through the model
output = model(input_image)

Running this code outputs:

0: 640x640 1 Banana, 70.7ms
Speed: 2.3ms preprocess, 70.7ms inference, 0.7ms postprocess per image at shape (1, 3, 640, 640)

I'd like to save Banana to a string.

output is an object of type list with one item in the list, which is an instance of Results from the ultralytics library. When trying to print output, I get the following:

ultralytics.engine.results.Results object with attributes:

boxes: ultralytics.engine.results.Boxes object
keypoints: None
masks: None
names: {0: 'Apple', 1: 'Banana', 2: 'Blueberry'}
obb: None
orig_img: array([[[255, 255, 255],
        [255, 255, 255],
        [255, 255, 255],
        ...,
        [255, 255, 255],
        [255, 255, 255],
        [255, 255, 255]],

       [[255, 255, 255],
        [255, 255, 255],
        [255, 255, 255],
        ...,
        [255, 255, 255],
        [255, 255, 255],
        [255, 255, 255]],

       [[255, 255, 255],
        [255, 255, 255],
        [255, 255, 255],
        ...,
        [255, 255, 255],
        [255, 255, 255],
        [255, 255, 255]],

       ...,

       [[255, 255, 255],
        [255, 255, 255],
        [255, 255, 255],
        ...,
        [255, 255, 255],
        [255, 255, 255],
        [255, 255, 255]],

       [[255, 255, 255],
        [255, 255, 255],
        [255, 255, 255],
        ...,
        [255, 255, 255],
        [255, 255, 255],
        [255, 255, 255]],

       [[255, 255, 255],
        [255, 255, 255],
        [255, 255, 255],
        ...,
        [255, 255, 255],
        [255, 255, 255],
        [255, 255, 255]]], dtype=uint8)
orig_shape: (225, 225)
path: 'fruit.jpeg'
probs: None
save_dir: 'runs/detect/predict'
speed: {'preprocess': 2.3441314697265625, 'inference': 70.66583633422852, 'postprocess': 0.6880760192871094}

Trying to print output.boxes gives me:

ultralytics.engine.results.Boxes object with attributes:

cls: tensor([0.])
conf: tensor([0.9832])
data: tensor([[ 63.5629,  44.9124, 153.1798, 183.1818,   0.9832,   0.0000]])
id: None
is_track: False
orig_shape: (225, 225)
shape: torch.Size([1, 6])
xywh: tensor([[108.3714, 114.0471,  89.6169, 138.2694]])
xywhn: tensor([[0.4817, 0.5069, 0.3983, 0.6145]])
xyxy: tensor([[ 63.5629,  44.9124, 153.1798, 183.1818]])
xyxyn: tensor([[0.2825, 0.1996, 0.6808, 0.8141]])

Nowhere in any of these outputs can I determine that the result was Banana.

How do I save the result (Banana) to a string?

Answer 1

Here you are:

print(output[0].names[np.array(output[0].boxes.cls)[0]])

There may be multiple classes detected, you may want to replace the last [0].