src/deep_learning/data/coco/dataset.py

import json
import math
from dataclasses import dataclass, field
from pathlib import Path

import keras
import numpy as np
import tensorflow as tf

from deep_learning.data.spec import SupervisedDataSource
from deep_learning.models.yolo import build_yolo_preprocessor, draw_prediction


@dataclass
class YoloDataSource(SupervisedDataSource):
    images_path: Path
    annotation_file: Path
    image_size: int = 448
    grid_size: int = 6
    batch_size: int = 32
    validation_batches: int = 500
    max_objects_per_image: int = 4
    example_count: int = 5
    example_output_dir: Path = Path("local/examples/yolo")

    _images_path: Path = field(init=False, repr=False)
    _annotation_file: Path = field(init=False, repr=False)
    _samples: list[dict] = field(init=False, repr=False, default_factory=list)

    def __post_init__(self):
        self._images_path = Path(self.images_path).expanduser()
        self._annotation_file = Path(self.annotation_file).expanduser()
        self.preprocessor = build_yolo_preprocessor(image_size=self.image_size)

    def sample_ds(self) -> tf.data.Dataset:
        samples = self._load_samples()
        return tf.data.Dataset.from_generator(
            lambda: (
                sample
                for sample in samples
            ),
            output_signature={
                "path": tf.TensorSpec(shape=(), dtype=tf.string),
                "boxes": tf.TensorSpec(shape=(None, 4), dtype=tf.float32),
                "labels": tf.TensorSpec(shape=(None,), dtype=tf.int32)
            }
        )

    def training_ds(self):
        dataset = self.sample_ds()
        dataset = dataset.filter(
            lambda sample: tf.shape(sample["boxes"])[0] <= self.max_objects_per_image
        )
        dataset = dataset.map(
            lambda sample: self._build_training_sample(sample),
            num_parallel_calls=tf.data.AUTOTUNE
        )
        dataset = dataset.batch(self.batch_size).prefetch(tf.data.AUTOTUNE)
        validation_ds = dataset.take(self.validation_batches)
        train_ds = dataset.skip(self.validation_batches)
        return train_ds, validation_ds

    def _build_training_sample(self, sample: dict) -> tuple[tf.Tensor, dict]:
        image = self._load_image(sample["path"])
        box_array, class_array = tf.numpy_function(
            self._build_label_arrays_numpy,
            [sample["boxes"], sample["labels"]],
            [tf.float32, tf.int32]
        )
        box_array.set_shape((self.grid_size, self.grid_size, 5))
        class_array.set_shape((self.grid_size, self.grid_size))
        return image, {"box": box_array, "class": class_array}

    def _load_samples(self) -> list[dict]:
        if self._samples:
            return self._samples

        with self._annotation_file.open("r", encoding="utf-8") as file:
            annotations = json.load(file)

        images = {image["id"]: image for image in annotations["images"]}
        metadata = {}
        for annotation in annotations["annotations"]:
            image_id = annotation["image_id"]
            if image_id not in metadata:
                image = images[image_id]
                metadata[image_id] = {
                    "path": str(self._images_path / image["file_name"]),
                    "boxes": [],
                    "labels": []
                }
            image = images[image_id]
            box = self.scale_box(annotation["bbox"], image["width"], image["height"])
            metadata[image_id]["boxes"].append(box)
            metadata[image_id]["labels"].append(annotation["category_id"])

        self._samples = [
            {
                "path": sample["path"],
                "boxes": np.asarray(sample["boxes"], dtype="float32"),
                "labels": np.asarray(sample["labels"], dtype="int32")
            }
            for sample in metadata.values()
        ]
        return self._samples

    def scale_box(self, box: list[float], width: int, height: int) -> list[float]:
        scale = 1.0 / max(width, height)
        x, y, w, h = [value * scale for value in box]
        if height > width:
            x += (height - width) * scale / 2
        if width > height:
            y += (width - height) * scale / 2
        return [x, y, w, h]

    def to_grid(self, box: list[float]) -> tuple[tuple[int, int], tuple[float, float, float, float]]:
        x, y, w, h = box
        center_x = (x + w / 2) * self.grid_size
        center_y = (y + h / 2) * self.grid_size
        index_x = int(center_x)
        index_y = int(center_y)
        return (index_x, index_y), (center_x - index_x, center_y - index_y, w, h)

    def _build_label_arrays_numpy(self, boxes: np.ndarray, labels: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
        box_array = np.zeros((self.grid_size, self.grid_size, 5), dtype="float32")
        class_array = np.zeros((self.grid_size, self.grid_size), dtype="int32")

        for box, label in zip(boxes, labels):
            x, y, w, h = box
            left = max(math.floor(x * self.grid_size), 0)
            right = min(math.ceil((x + w) * self.grid_size), self.grid_size)
            bottom = max(math.floor(y * self.grid_size), 0)
            top = min(math.ceil((y + h) * self.grid_size), self.grid_size)
            class_array[bottom:top, left:right] = label

        for box, label in zip(boxes, labels):
            (index_x, index_y), grid_box = self.to_grid(box.tolist())
            index_x = min(max(index_x, 0), self.grid_size - 1)
            index_y = min(max(index_y, 0), self.grid_size - 1)
            box_array[index_y, index_x] = [*grid_box, 1.0]
            class_array[index_y, index_x] = label

        return box_array, class_array

    def _load_image(self, path: tf.Tensor) -> tf.Tensor:
        image = tf.io.read_file(path)
        image = tf.image.decode_jpeg(image, channels=3)
        image = tf.expand_dims(image, axis=0)
        image = self.preprocessor(image)
        image = tf.squeeze(image, axis=0)
        return tf.cast(image, tf.float32)

    def test_examples(self, model: keras.Model) -> None:
        self.example_output_dir.mkdir(parents=True, exist_ok=True)
        examples_ds = self.sample_ds().take(self.example_count)
        for index, sample in enumerate(examples_ds, start=1):
            input_path = Path(sample["path"].numpy().decode("utf-8"))
            image = self._load_example_image(sample["path"])
            image_array = image.numpy()
            predictions = model.predict(np.expand_dims(image_array, axis=0), verbose=0)
            boxes = np.asarray(predictions["box"][0])
            classes = np.argmax(np.asarray(predictions["class"][0]), axis=-1)
            result_image = draw_prediction(
                image_array,
                boxes,
                classes,
                self.grid_size,
                cutoff=0.2
            )
            prediction_path = self.example_output_dir / f"example_{index:03d}_prediction.png"
            result_image.save(prediction_path)
            print(
                f"样例 {index}: "
                f"输入={input_path}, "
                f"检测结果={prediction_path}"
            )

    @staticmethod
    def _load_example_image(path: tf.Tensor) -> tf.Tensor:
        image = tf.io.read_file(path)
        image = tf.image.decode_jpeg(image, channels=3)
        return tf.cast(image, tf.uint8)