PriceEstimation/classify.py

from ultralytics import YOLO


def predict_class(model_path, img_path, conf=0.7):
    # Load the model
    model = YOLO(model_path)
    # Get the results for the given image
    results_list = model(img_path, conf=conf)
    
    # Ensure results_list is not empty and access the first result
    if results_list:
        results = results_list[0]
        # Access the probabilities
        probs = results.probs
        if probs is not None:
            # Find the class index with the highest probability (top1)
            class_idx = probs.top1
            class_name = results.names[class_idx]
            return class_name
        else:
            return "No probabilities found in the results."
    else:
        return "No results returned by the model."