courtvision.trainer

log_wb_image_and_bbox(images, preds, targets, logger, global_step)

summary

Parameters:

Name	Type	Description	Default
images	torch.tensor	description	required
preds	list[dict[str, torch.Tensor]]	description	required
targets	list[dict[str, torch.Tensor]]	description	required
logger	wandb.sdk.wandb_run.Run	description	required
global_step	int	description	required

Source code in courtvision/trainer.py

def log_wb_image_and_bbox(
    images: torch.tensor,
    preds: list[dict[str, torch.Tensor]],
    targets: list[dict[str, torch.Tensor]],
    logger: wandb.sdk.wandb_run.Run,
    global_step: int,
):
    """_summary_

    Args:
        images (torch.tensor): _description_
        preds (list[dict[str, torch.Tensor]]): _description_
        targets (list[dict[str, torch.Tensor]]): _description_
        logger (wandb.sdk.wandb_run.Run): _description_
        global_step (int): _description_
    """
    images_to_log = []
    for image, pred, target in zip(images, preds, targets):
        image_height, image_width = image.shape[-2:]
        images_to_log.append(
            wandb.Image(
                image.permute(1, 2, 0).cpu().numpy(),
                boxes={
                    "predictions": {
                        "box_data": [
                            {
                                "position": {
                                    "minX": float(x_min) / image_width,
                                    "maxX": float(x_max) / image_width,
                                    "minY": float(y_min) / image_height,
                                    "maxY": float(y_max) / image_height,
                                },
                                "class_id": 1,
                                "box_caption": "ball",
                                "scores": {"score": float(score)},
                            }
                            for (x_min, y_min, x_max, y_max), score in zip(
                                pred["boxes"].cpu().numpy(),
                                pred["scores"].cpu().numpy(),
                                strict=True,
                            )
                        ]
                    },
                    "targets": {
                        "box_data": [
                            {
                                "position": {
                                    "minX": float(x_min) / image_width,
                                    "maxX": float(x_max) / image_width,
                                    "minY": float(y_min) / image_height,
                                    "maxY": float(y_max) / image_height,
                                },
                                "class_id": 1,
                                "box_caption": "ball",
                            }
                            for x_min, y_min, x_max, y_max in target["boxes"]
                            .cpu()
                            .numpy()
                        ],
                    },
                },
            )
        )

    logger.log({"image": images_to_log})