Spaces:
Sleeping
Sleeping
File size: 2,899 Bytes
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 | from dataclasses import dataclass
from pathlib import Path
import keras
import keras_hub
from keras import layers
from .losses import box_loss
from deep_learning.models.spec import ModelArtifact, SupervisedModelBuilder
def build_yolo_preprocessor(image_size: int = 448):
inputs = keras.Input(shape=(None, None, 3), dtype="uint8")
x = layers.Resizing(
image_size,
image_size,
interpolation="bicubic",
crop_to_aspect_ratio=True
)(inputs)
x = layers.Rescaling(
scale=[0.017124753831663668, 0.01750700280112045, 0.017429193899782133],
offset=[-2.1179039301310043, -2.0357142857142856, -1.8044444444444445]
)(x)
return keras.Model(inputs, x, name="yolo_preprocessor")
@dataclass
class YoloModelBuilder(SupervisedModelBuilder):
image_size: int
grid_size: int
num_labels: int
backbone_preset: str = "resnet_50_imagenet"
def build_training_artifact(self) -> ModelArtifact:
backbone = keras_hub.models.Backbone.from_preset(self.backbone_preset)
inputs = keras.Input(shape=(self.image_size, self.image_size, 3))
x = backbone(inputs)
x = layers.Conv2D(512, (3, 3), strides=(2, 2))(x)
x = layers.Flatten()(x)
x = layers.Dense(2048, activation="relu", kernel_initializer="glorot_normal")(x)
x = layers.Dropout(0.5)(x)
x = layers.Dense(self.grid_size * self.grid_size * (self.num_labels + 5))(x)
x = layers.Reshape((self.grid_size, self.grid_size, self.num_labels + 5))(x)
box_predictions = x[..., :5]
class_predictions = layers.Activation("softmax")(x[..., 5:])
outputs = {"box": box_predictions, "class": class_predictions}
model = keras.Model(inputs, outputs, name="yolo")
return ModelArtifact(model=model)
def convert_to_inference_artifact(
self,
training_artifact: ModelArtifact
) -> ModelArtifact:
inputs = keras.Input(shape=(None, None, 3), dtype="uint8")
preprocessor = build_yolo_preprocessor(image_size=self.image_size)
x = preprocessor(inputs)
outputs = training_artifact.model(x)
inference_model = keras.Model(inputs, outputs, name="yolo_inference")
return ModelArtifact(model=inference_model)
def load_inference_artifact(self, model_path: Path) -> ModelArtifact:
model = keras.models.load_model(
str(model_path),
custom_objects=self._custom_objects()
)
return ModelArtifact(model=model)
def _custom_objects(self) -> dict:
return {
"box_loss": box_loss
}
def compile_training_model(self, model: keras.Model) -> None:
model.compile(
optimizer=keras.optimizers.Adam(2e-4),
loss={
"box": box_loss,
"class": "sparse_categorical_crossentropy"
}
)
|