Almost every vision project we get called into starts as "we need to classify these images" and turns out to be a detection or segmentation problem within a week. Counting pallets in a warehouse photo, finding the serial-number plate before you OCR it, measuring the area of a defect on a weld: none of those are answered by a single label per image. Our vision transfer-learning tutorial covers the classification case. This one covers the two harder heads - object detection (boxes) and segmentation (masks) - on the TensorFlow and Keras 3 stack, including the export path to LiteRT for on-device use.
What to reach for
| Task | Reasonable starting point |
|---|---|
| Fixed set of classes, plenty of labels | Fine-tune a KerasHub detector (RetinaNet / YOLO-family preset) on your data |
| Fixed classes, very few labels | Freeze a pretrained backbone, train only the heads; augment hard |
| Classes you cannot enumerate up front | Open-vocabulary detection with text prompts, then distill into a fixed-class model |
| Cutouts / area measurement, few labels | Promptable segmentation (Segment Anything-style) to pre-label, then train a small U-Net for production |
The last row matters commercially: promptable segmentation models are excellent annotation accelerators and usually too heavy to serve on a device. The winning pattern is "big model labels, small model ships."
Setup
pip install -U keras keras-hub tensorflow
import os
os.environ["KERAS_BACKEND"] = "tensorflow"
import keras, keras_hub, tensorflow as tf
keras.mixed_precision.set_global_policy("mixed_float16")
Check which detection and segmentation presets your installed keras-hub version actually ships - the model zoo moves quickly:
for name in dir(keras_hub.models):
if "Detect" in name or "Segment" in name:
print(name)
Get the labels into a sane format
Detection data is where these projects die. Pick one canonical format, convert once, and validate it. We standardise on a tf.data pipeline yielding a dict per example:
# boxes in "xyxy" absolute pixels, classes as int32 ids
def decode(record):
image = tf.io.decode_jpeg(tf.io.read_file(record["path"]), channels=3)
return {
"images": tf.cast(image, tf.float32),
"bounding_boxes": {
"boxes": record["boxes"], # (num_objects, 4)
"classes": record["classes"], # (num_objects,)
},
}
train_ds = (
tf.data.Dataset.from_tensor_slices(records)
.map(decode, num_parallel_calls=tf.data.AUTOTUNE)
.ragged_batch(8)
.prefetch(tf.data.AUTOTUNE)
)
Three checks worth automating before you train anything:
- Draw the boxes. Render 50 random training samples with their boxes burned in and look at them. Coordinate-order bugs (
xyxyvsyxyxvsxywh, absolute vs relative) are the single most common defect and they are invisible in the loss curve for the first few epochs. - Box sanity. No zero-area boxes, no coordinates outside the image, no class ids outside your label map.
- Class balance and object size histogram. If half your objects are under 16 px on the long side, you need higher input resolution or tiling, not a better model.
Augmentation that respects boxes
Geometric augmentation has to transform the labels too. Use the Keras layers that are box-aware and keep the format explicit:
augment = keras.Sequential([
keras.layers.RandomFlip("horizontal", bounding_box_format="xyxy"),
keras.layers.RandomZoom(0.2, bounding_box_format="xyxy"),
])
Colour jitter (brightness, contrast, saturation) is free and safe because it never moves a box. Heavy rotation often is not safe: axis-aligned boxes around a rotated object get looser, which teaches the model sloppy localisation. If your product photos are always upright, do not rotate.
Fine-tune a detector
detector = keras_hub.models.ImageObjectDetector.from_preset(
"retinanet_resnet50_fpn_coco",
num_classes=len(LABELS),
bounding_box_format="xyxy",
)
detector.compile(
optimizer=keras.optimizers.AdamW(learning_rate=1e-4, global_clipnorm=10.0),
loss=None, # preset detectors define their own box + classification losses
)
detector.fit(train_ds, validation_data=val_ds, epochs=30,
callbacks=[keras.callbacks.ReduceLROnPlateau(patience=3),
keras.callbacks.EarlyStopping(patience=6, restore_best_weights=True)])
Notes from real runs:
- Gradient clipping is not optional. Detection losses spike;
global_clipnormkeeps a bad batch from wrecking the run. - Two-stage schedule. Freeze the backbone for the first few epochs so the randomly initialised heads settle, then unfreeze at a 10x lower learning rate.
- Resolution beats architecture for small objects. Going from 640 to 960 px input usually helps more than swapping the backbone, at roughly quadratic cost.
Score it the way your business cares about
mAP over IoU 0.5:0.95 is the paper metric. It is rarely the metric your client is paid on. Compute mAP for monitoring, but also report the operating point you will actually deploy:
metric = keras_hub.metrics.MeanAveragePrecision(bounding_box_format="xyxy") # or a COCO eval harness
Then fix a confidence threshold and report precision and recall per class at that threshold, plus the cost of each error type. "We miss 4% of defects and raise 11% false alarms at threshold 0.45" is a sentence a plant manager can make a decision with; "mAP 0.612" is not. Tune the threshold and the NMS IoU on the validation set, never on the test set, and write the chosen values into the model card (see our explainability tutorial).
Segmentation: pre-label big, ship small
For masks, use a promptable segmentation model to generate candidate masks from cheap prompts (a click or a box per object), have a human accept or fix them, and train a small production model on the result:
sam = keras_hub.models.SAMImageSegmenter.from_preset("sam_base_sa1b")
masks = sam.predict({
"images": image_batch,
"points": point_batch, # one click per object
"labels": label_batch, # 1 = foreground
})
Then a compact U-Net-style Keras model trained on those masks with Dice plus binary cross-entropy loss:
model.compile(
optimizer=keras.optimizers.AdamW(1e-3),
loss=lambda y, p: keras.losses.binary_crossentropy(y, p) + dice_loss(y, p),
metrics=[keras.metrics.BinaryIoU(target_class_ids=[1])],
)
Report IoU per class and boundary error, not pixel accuracy - with a small foreground, predicting all background scores 97% pixel accuracy and is worthless.
Export and serve
Server-side, export a SavedModel with an explicit signature that takes encoded bytes so clients never have to match your preprocessing, exactly as in our TF Serving tutorial:
@tf.function(input_signature=[tf.TensorSpec([None], tf.string)])
def serve_jpeg(encoded):
images = tf.map_fn(lambda b: tf.image.resize(
tf.io.decode_jpeg(b, channels=3), (640, 640)), encoded,
fn_output_signature=tf.float32)
return detector(images)
On-device, convert to LiteRT and keep the post-processing honest:
converter = tf.lite.TFLiteConverter.from_saved_model("export/detector/1")
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.representative_dataset = rep_ds # int8 needs real images
tflite_model = converter.convert()
Two traps we hit on nearly every edge detection project:
- NMS on device. Non-max suppression either becomes a select-op the runtime must support, or it lives in your application code. Decide deliberately; do not discover it at integration time.
- Quantisation hurts localisation more than classification. Always re-measure mAP after int8 conversion on real device output, not just the class scores. See Migrating from TensorFlow Lite to LiteRT and Shrinking TensorFlow Models for the conversion details.
A realistic project shape
For a fixed-class industrial detection task with a few thousand labelled images, we plan roughly: one week to nail the label format and audit the annotations, one week of training and threshold work, one week for export, device benchmarking, and the monitoring hooks. The modelling is the short part. Data quality and the deployment contract are where the schedule actually goes.
If you are staring at a folder of images and an unlabelled backlog, that is a normal place to start from - get in touch and we will scope the labelling strategy and the model together, or see our services for how these engagements are structured.