| from typing import Any, Optional | |
| import numpy as np | |
| from PIL import Image | |
| from tqdm.auto import tqdm | |
| from ultralytics.engine.results import Results | |
| from ultralytics.models.yolo.pose import PosePredictor | |
| from ultralytics.utils.plotting import Annotator | |
| import wandb | |
| from wandb.integration.ultralytics.bbox_utils import ( | |
| get_boxes, | |
| get_ground_truth_bbox_annotations, | |
| ) | |
| def annotate_keypoint_results(result: Results, visualize_skeleton: bool): | |
| annotator = Annotator(np.ascontiguousarray(result.orig_img[:, :, ::-1])) | |
| key_points = result.keypoints.data.numpy() | |
| for idx in range(key_points.shape[0]): | |
| annotator.kpts(key_points[idx], kpt_line=visualize_skeleton) | |
| return annotator.im | |
| def annotate_keypoint_batch(image_path: str, keypoints: Any, visualize_skeleton: bool): | |
| with Image.open(image_path) as original_image: | |
| original_image = np.ascontiguousarray(original_image) | |
| annotator = Annotator(original_image) | |
| annotator.kpts(keypoints.numpy(), kpt_line=visualize_skeleton) | |
| return annotator.im | |
| def plot_pose_predictions( | |
| result: Results, | |
| model_name: str, | |
| visualize_skeleton: bool, | |
| table: Optional[wandb.Table] = None, | |
| ): | |
| result = result.to("cpu") | |
| boxes, mean_confidence_map = get_boxes(result) | |
| annotated_image = annotate_keypoint_results(result, visualize_skeleton) | |
| prediction_image = wandb.Image(annotated_image, boxes=boxes) | |
| table_row = [ | |
| model_name, | |
| prediction_image, | |
| len(boxes["predictions"]["box_data"]), | |
| mean_confidence_map, | |
| result.speed, | |
| ] | |
| if table is not None: | |
| table.add_data(*table_row) | |
| return table | |
| return table_row | |
| def plot_pose_validation_results( | |
| dataloader, | |
| class_label_map, | |
| model_name: str, | |
| predictor: PosePredictor, | |
| visualize_skeleton: bool, | |
| table: wandb.Table, | |
| max_validation_batches: int, | |
| epoch: Optional[int] = None, | |
| ) -> wandb.Table: | |
| data_idx = 0 | |
| num_dataloader_batches = len(dataloader.dataset) // dataloader.batch_size | |
| max_validation_batches = min(max_validation_batches, num_dataloader_batches) | |
| for batch_idx, batch in enumerate(dataloader): | |
| prediction_results = predictor(batch["im_file"]) | |
| progress_bar_result_iterable = tqdm( | |
| enumerate(prediction_results), | |
| total=len(prediction_results), | |
| desc=f"Generating Visualizations for batch-{batch_idx + 1}/{max_validation_batches}", | |
| ) | |
| for img_idx, prediction_result in progress_bar_result_iterable: | |
| prediction_result = prediction_result.to("cpu") | |
| table_row = plot_pose_predictions( | |
| prediction_result, model_name, visualize_skeleton | |
| ) | |
| ground_truth_image = wandb.Image( | |
| annotate_keypoint_batch( | |
| batch["im_file"][img_idx], | |
| batch["keypoints"][img_idx], | |
| visualize_skeleton, | |
| ), | |
| boxes={ | |
| "ground-truth": { | |
| "box_data": get_ground_truth_bbox_annotations( | |
| img_idx, batch["im_file"][img_idx], batch, class_label_map | |
| ), | |
| "class_labels": class_label_map, | |
| }, | |
| }, | |
| ) | |
| table_row = [data_idx, batch_idx, ground_truth_image] + table_row[1:] | |
| table_row = [epoch] + table_row if epoch is not None else table_row | |
| table_row = [model_name] + table_row | |
| table.add_data(*table_row) | |
| data_idx += 1 | |
| if batch_idx + 1 == max_validation_batches: | |
| break | |
| return table | |