from roboflow import Roboflow
from ultralytics import YOLO
import cv2
import matplotlib.pyplot as plt

# Step 2: Load the dataset
rf = Roboflow(api_key="YOUR_ROBOFLOW_API_KEY")
project = rf.workspace().project("YOUR_PROJECT_NAME")
dataset = project.version("YOUR_DATASET_VERSION").download("yolov8")

# Step 3: Define the YOLOv8 model
model = YOLO('yolov8n.yaml')

# Step 4: Train the model
model.train(data=f"{dataset.location}/data.yaml", epochs=50, imgsz=640, batch=16)

# Step 5: Evaluate the model
results = model.val()
print(results)

# Step 6: Test the model
# Load and preprocess the image
img_path = 'path/to/your/test/image.jpg'
img = cv2.imread(img_path)
img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

# Run predictions
results = model(img_rgb)

# Display the results
plt.imshow(results.render()[0])
plt.show()