Source code for tmas.src.models.yolo

# tmas/tmas/models/yolo.py
import os
from ultralytics import YOLO
import os

# Get the directory of the current file
[docs] current_dir = os.path.dirname(__file__)
# print(f'YOLO file curent path: {current_dir} ') # Construct the full path to the model file
[docs] weights = os.path.join(current_dir, 'best_model_yolo.pt')
# print(f'YOLO weight path: {weights} ')
[docs] class YOLOv8: def __init__(self, weights= weights): if not os.path.exists(weights): raise FileNotFoundError(f"Model file {weights} not found. Please ensure the model file is in the correct path.") # Load the YOLOv8 model
[docs] self.model = YOLO(weights)
[docs] def predict(self, img): # Predict using the loaded model return self.model(img)