To perform object classification on an image file using Python, we can use the open source pre-trained YOLO model from Ultralytics.
First, install the library using:
$ pip install ultralytics
For example, assume we have an image of a tractor in a local file tractor.jpeg under images/.
Note that we can also run the model from the command line using:
$ yolo predict source='images/tractor.jpeg'
In Python, we need to extract the result from all of the model output, which requires a bit more code.
The model’s predict function will return a list of results with probability values, as well as a list of all labels.
The code below will extract the highest probability label and print it.
from ultralytics import YOLO model = YOLO("yolov8n-cls.pt") # Path to an image file assumed to exist. results = model.predict("images/tractor.jpeg") # Overall results is a list. result = results probabilities = result.probs # Top1 is the most likely result. topLabelNumber = probabilities.top1 # Now find the label name for that label number. allNames = result.names for labelNumber, label in allNames.items(): if labelNumber == topLabelNumber: resultLabel = label print("Classification result:") print(resultLabel)