detection_yolo.py

import wandb
from wandb.integration.ultralytics import add_wandb_callback
from ultralytics import YOLO

model = YOLO("./runs/detect/tune/weights/best.pt")

add_wandb_callback(model, enable_model_checkpointing=True)

wandb.init(
    project="my-yolo-project", 
    name="running-cocoa", 
    config={
        "epochs": 100,
        "imgsz": 640,
    }
)

def on_train_epoch_end(trainer):
   try:
    print(f"Callback triggered at epoch {trainer.epoch}")
    wandb.log({
        "mAP50": trainer.metrics.get('metrics/mAP50(B)', 0),
        "mAP50-95": trainer.metrics.get('metrics/mAP50-95(B)', 0),
        "precision": trainer.metrics.get('metrics/precision(B)', 0),
        "recall": trainer.metrics.get('metrics/recall(B)', 0),  
        "learning_rate": trainer.optimizer.param_groups[0]['lr'],
    }, step=trainer.epoch)
   except:
    print(f"Error in on_train_epoch_end: {e}")

model.add_callback("on_train_epoch_end", on_train_epoch_end)

train_results = model.train(
    data="data.yaml",
    epochs=100,
    imgsz=640,
    project="runs/final_train",
    name="cocoa-detection-run",
    verbose=True,
    workers=0
)

inference_results = model(
    [
        "./safe_copy/50084_1.jpg",
        "./safe_copy/50084_2.jpg",
        "./safe_copy/B_2.jpg",
    ]
)

def log_predictions(img_paths, results):
    for img_path, result in zip(img_paths, results):
        boxes = result.boxes
        img = wandb.Image(
            img_path,
            boxes={
                "predictions": {
                    "box_data": [{
                        "position": {
                            "minX": float(box.xyxy[0][0]),
                            "minY": float(box.xyxy[0][1]),
                            "maxX": float(box.xyxy[0][2]),
                            "maxY": float(box.xyxy[0][3])
                        },
                        "class_id": int(box.cls),
                        "confidence": float(box.conf)
                    } for box in boxes],
                    "class_labels": model.names
                }
            }
        )
        wandb.log({"predictions": img})

log_predictions(
    ["./safe_copy/50084_1.jpg",
     "./safe_copy/50084_2.jpg",
     "./safe_copy/B_2.jpg",],
    inference_results
)

wandb.finish()