Spaces:
Sleeping
Sleeping
File size: 7,370 Bytes
14f6839 07cb7d3 14f6839 07cb7d3 14f6839 07cb7d3 14f6839 07cb7d3 14f6839 07cb7d3 14f6839 07cb7d3 14f6839 07cb7d3 14f6839 07cb7d3 14f6839 07cb7d3 14f6839 07cb7d3 14f6839 07cb7d3 14f6839 07cb7d3 14f6839 07cb7d3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 | 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)
|