soumik12345 / wandb-addons

Weights & Biases Addons is a repository consisting of additional unitilities and community contributions for supercharging your Weights & Biases workflows.
https://geekyrakshit.dev/wandb-addons/
23 stars 7 forks source link

feat: Callback for `ultralytics` for logging bounding boxes in `wandb.Table` #35

Closed soumik12345 closed 1 year ago

soumik12345 commented 1 year ago

Fixes https://github.com/soumik12345/wandb-addons/issues/34

Callback for logging model checkpoints, predictions, and ground-truth annotations with interactive overlays for bounding boxes to Weights & Biases Tables during training, validation, and prediction for a ultratytics workflow.

Inspired highly by https://github.com/wandb/wandb/pull/5168

Sample Result: https://wandb.ai/geekyrakshit/YOLOv8/reports/Ultralytics-Integration-Demo--Vmlldzo0Nzk5OTEz

Sample Usage:

git clone https://github.com/soumik12345/wandb-addons
pip install ./wandb-addons[yolo]
from ultralytics.yolo.engine.model import YOLO

import wandb
from wandb_addons.ultralytics import add_wandb_callback

# initialize wandb run
wandb.init(project="YOLOv8")

# initialize YOLO model
model = YOLO("yolov8n.pt")

# add wandb callback
add_wandb_callback(model, max_validation_batches=2, enable_model_checkpointing=True)

# train
model.train(data="coco128.yaml", epochs=5, imgsz=640)

# validate
model.val()

# perform inference
model(['img1.jpeg', 'img2.jpeg'])