Spaces:
Sleeping
Sleeping
| 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." |