feat: Enable custom loss callbacks for integration with wandb/mlflow/tensorboard

[yuguerten]

use case:

import wandb
from cellpose import models
from pathlib import Path

# Initialize experiment logging (e.g., W&B)
wandb.init(
    project="cellpose-bacteria-segmentation",
    config={
        "learning_rate": 0.0001,
        "epochs": 100,
        "batch_size": 8,
        "model_type": "cyto"
    }
)

# Define callback to log metrics to wandb
def wandb_callback(epoch, train_loss, test_loss):
    """Log training metrics to Weights & Biases"""
    metrics = {"epoch": epoch, "train_loss": train_loss}
    if test_loss is not None:
        metrics["test_loss"] = test_loss
    wandb.log(metrics)

# Load model and train
model = models.CellposeModel(gpu=True)

model_path = train_seg(
    net=model.net,
    train_files=train_files,
    train_labels_files=train_labels_files,
    test_files=test_files,
    test_labels_files=test_labels_files,
    n_epochs=100,
    learning_rate=0.0001,
    loss_callback=wandb_callback,  # Add wandb logging
    return_loss_arrays=False
)

# Save model artifact
wandb.log_artifact(wandb.Artifact('model', type='model').add_file(model_path))
wandb.finish()