<commit_before># Copyright 2022 The KerasCV Authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import six
import tensorflow as tf
from keras.callbacks import Callback

from keras_cv import bounding_box
from keras_cv.metrics.coco import compute_pycoco_metrics


class PyCOCOCallback(Callback):
    def __init__(self, validation_data, bounding_box_format, **kwargs):
        self.model = None
        self.val_data = validation_data
        self.bounding_box_format = bounding_box_format
        super().__init__(**kwargs)

    def on_epoch_end(self, epoch, logs=None):
        logs = logs or {}

        gt = {}
        preds = {}

        for i, batch in enumerate(self.val_data):
            gt_i, preds_i = self._eval_batch(batch, i)

            for k, v in six.iteritems(preds_i):
                if k not in preds:
                    preds[k] = [v]
                else:
                    preds[k].append(v)

            for k, v in six.iteritems(gt_i):
                if k not in gt:
                    gt[k] = [v]
                else:
                    gt[k].append(v)

        metrics = compute_pycoco_metrics(gt, preds)
        # Mark these as validation metrics by prepending a val_ prefix
        metrics = {"val_" + name: val for name, val in metrics.items()}

        logs.update(metrics)

    def _eval_batch(self, batch, index):
        images, y = batch
        gt_boxes = y["gt_boxes"]
        gt_classes = y["gt_classes"]
        batch_size = images.shape[0]
        height = images.shape[1]
        width = images.shape[2]

        gt_boxes = bounding_box.convert_format(
            gt_boxes, source=self.bounding_box_format, target="yxyx"
        )

        source_ids = tf.strings.join(
            [
                tf.strings.as_string(tf.tile(tf.constant([index]), [batch_size])),
                tf.strings.as_string(
                    tf.linspace(1, batch_size, batch_size), precision=0
                ),
            ],
            separator="/",
        )

        ground_truth = {}
        ground_truth["source_id"] = source_ids
        ground_truth["height"] = tf.tile(tf.constant([height]), [batch_size])
        ground_truth["width"] = tf.tile(tf.constant([width]), [batch_size])

        num_dets = gt_classes.get_shape().as_list()[1]
        ground_truth["num_detections"] = tf.tile(tf.constant([num_dets]), [batch_size])
        ground_truth["boxes"] = gt_boxes
        ground_truth["classes"] = gt_classes

        y_pred = self.model.predict(images)
        y_pred = bounding_box.convert_format(
            y_pred, source=self.bounding_box_format, target="yxyx"
        )

        predictions = {}
        predictions["num_detections"] = y_pred.row_lengths()
        y_pred = y_pred.to_tensor(-1)

        predictions["source_id"] = source_ids
        predictions["detection_boxes"] = y_pred[:, :, :4]
        predictions["detection_classes"] = y_pred[:, :, 4]
        predictions["detection_scores"] = y_pred[:, :, 5]

        return ground_truth, predictions<commit_msg>[Object Detection] Speed up the PyCOCO callback by only calling model.predict once (#1026)<commit_after># Copyright 2022 The KerasCV Authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import tensorflow as tf
from keras.callbacks import Callback

from keras_cv import bounding_box
from keras_cv.metrics.coco import compute_pycoco_metrics


class PyCOCOCallback(Callback):
    def __init__(self, validation_data, bounding_box_format, cache=True, **kwargs):
        """Creates a callback to evaluate PyCOCO metrics on a validation dataset.

        Args:
            validation_data: a tf.data.Dataset containing validation data. Entries
                should have the form ```(images, {"gt_boxes": boxes,
                "gt_classes": classes})```.
            bounding_box_format: the KerasCV bounding box format used in the
                validation dataset (e.g. "xywh")
            cache: whether the callback should cache the dataset between iterations.
                Note that if the validation dataset has shuffling of any kind
                (e.g from `shuffle_files=True` in a call to TFDS.load or a call
                to tf.data.Dataset.shuffle() with `reshuffle_each_iteration=True`),
                you **must** cache the dataset to preserve iteration order. This
                will store your entire dataset in main memory, so for large datasets
                consider avoiding shuffle operations and passing `cache=False`.
        """
        self.model = None
        self.val_data = validation_data
        if cache:
            # We cache the dataset to preserve a consistent iteration order.
            self.val_data = self.val_data.cache()
        self.bounding_box_format = bounding_box_format
        super().__init__(**kwargs)

    def on_epoch_end(self, epoch, logs=None):
        logs = logs or {}

        def images_only(images, boxes):
            return images

        def boxes_only(images, boxes):
            return boxes

        images_only_ds = self.val_data.map(images_only)
        y_pred = self.model.predict(images_only_ds)

        gt = [boxes for boxes in self.val_data.map(boxes_only)]
        gt_boxes = tf.concat(
            [tf.RaggedTensor.from_tensor(boxes["gt_boxes"]) for boxes in gt], axis=0
        )
        gt_classes = tf.concat(
            [tf.RaggedTensor.from_tensor(boxes["gt_classes"]) for boxes in gt], axis=0
        )

        first_image_batch = next(iter(images_only_ds))
        height = first_image_batch.shape[1]
        width = first_image_batch.shape[2]
        total_images = gt_boxes.shape[0]

        gt_boxes = bounding_box.convert_format(
            gt_boxes, source=self.bounding_box_format, target="yxyx"
        )

        source_ids = tf.strings.as_string(
            tf.linspace(1, total_images, total_images), precision=0
        )

        ground_truth = {}
        ground_truth["source_id"] = [source_ids]
        ground_truth["height"] = [tf.tile(tf.constant([height]), [total_images])]
        ground_truth["width"] = [tf.tile(tf.constant([width]), [total_images])]

        ground_truth["num_detections"] = [gt_boxes.row_lengths(axis=1)]
        ground_truth["boxes"] = [gt_boxes.to_tensor(-1)]
        ground_truth["classes"] = [gt_classes.to_tensor(-1)]

        y_pred = bounding_box.convert_format(
            y_pred, source=self.bounding_box_format, target="yxyx"
        )

        predictions = {}
        predictions["num_detections"] = [y_pred.row_lengths()]
        y_pred = y_pred.to_tensor(-1)

        predictions["source_id"] = [source_ids]
        predictions["detection_boxes"] = [y_pred[:, :, :4]]
        predictions["detection_classes"] = [y_pred[:, :, 4]]
        predictions["detection_scores"] = [y_pred[:, :, 5]]

        metrics = compute_pycoco_metrics(ground_truth, predictions)
        # Mark these as validation metrics by prepending a val_ prefix
        metrics = {"val_" + name: val for name, val in metrics.items()}

        logs.update(metrics)
{'fraction_non_alphanumeric': 0.1004121393780442, 'fraction_numerical': 0.005495191707256151, 'mean_word_length': 2.5718108831400537, 'pattern_counts': {'":': 2, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 2, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2255469', 'n_tokens_mistral': 2391, 'n_tokens_neox': 2290, 'n_words': 628}