Ditching the Bloat: Building a Native Keras YOLOv8 Edge Object Detector (TFLite & CoreML)

If you have ever tried to deploy an object detection model to a mobile phone or a Raspberry Pi, you know the pain. You train a massive, state-of-the-art model using a heavy third-party framework, only to realize that converting it to TensorFlow Lite or CoreML requires the digital equivalent of black magic. Ops are unsupported, dynamic tensor shapes crash the compiler, and suddenly your 99% accurate model is completely useless on a mobile device.

And then there is the elephant in the room: Licensing.

Frameworks like Ultralytics are incredibly user-friendly, but their YOLOv8 implementation is bound by an AGPL-3.0 license. This means if you want to put your model into a closed-source, commercial mobile app, you are legally required to buy a rather expensive Enterprise license.

But here is the industry secret: the YOLOv8 architecture itself is just math. By building it natively using KerasCV — which operates under the extremely permissive, commercially friendly Apache 2.0 license — you get the exact same state-of-the-art performance, zero dependency bloat, and absolutely zero licensing fees.

In this guide, we are going to build, train, and convert a YOLOv8-Extra Small model for mask detection purely in Keras. No wrappers. Just clean, static graphs perfectly optimized for TFLite (Android/NPUs) and CoreML (iOS/Apple Neural Engine).

Let’s build it.

Step 1: Grabbing the Data

We need data, and we need it in a format that is easy to parse without relying on external libraries. We are using a mask-wearing dataset hosted on Roboflow, exported in Pascal VOC format. Why VOC? Because it generates simple .xml files that standard Python can read in its sleep.

!pip install roboflow keras_cv

from roboflow import Roboflow
rf = Roboflow(api_key="YOUR_API_KEY")
project = rf.workspace("joseph-nelson").project("mask-wearing")
version = project.version(19)
dataset = version.download("voc")

Next, we write a native pipeline to parse these XMLs and extract the bounding boxes (in xyxy format) and class labels.

import tensorflow as tf
import keras_cv
import os
import xml.etree.ElementTree as ET

CLASSES = ["mask", "no-mask"]
CLASS_MAPPING = {name: idx for idx, name in enumerate(CLASSES)}
BATCH_SIZE = 8

def parse_voc_annotation(xml_file):
    tree = ET.parse(xml_file)
    root = tree.getroot()
    boxes, classes = [], []

    for obj in root.findall('object'):
        class_name = obj.find('name').text
        if class_name not in CLASS_MAPPING: continue

        classes.append(CLASS_MAPPING[class_name])
        bndbox = obj.find('bndbox')
        
        boxes.append([
            float(bndbox.find('xmin').text), float(bndbox.find('ymin').text),
            float(bndbox.find('xmax').text), float(bndbox.find('ymax').text)
        ])
    return boxes, classes

def load_data(image_dir):
    image_paths, all_boxes, all_classes = [], [], []
    for img_name in os.listdir(image_dir):
        if not img_name.endswith('.jpg'): continue
        
        img_path = os.path.join(image_dir, img_name)
        xml_path = os.path.join(image_dir, img_name.replace('.jpg', '.xml'))
        
        if not os.path.exists(xml_path): continue
            
        boxes, classes = parse_voc_annotation(xml_path)
        image_paths.append(img_path)
        all_boxes.append(boxes)
        all_classes.append(classes)
        
    return image_paths, all_boxes, all_classes

train_img_paths, train_boxes, train_classes = load_data("Mask-Wearing-19/train")
valid_img_paths, valid_boxes, valid_classes = load_data("Mask-Wearing-19/valid")

Step 2: The XLA Boss Fight (Ragged vs. Dense Tensors)

Here is where most edge pipelines die. Images have a varying number of objects. One image might have 2 faces, another might have 10. Naturally, TensorFlow handles this using RaggedTensors (tensors with jagged edges).

The problem? The XLA compiler — the engine that makes model training ridiculously fast on GPUs — despises ragged tensors. It needs perfectly square, predictable blocks of memory. If we feed it ragged tensors, it will panic and crash.

To fix this, we pad our missing bounding boxes with -1.0. KerasCV is incredibly smart; its loss functions are explicitly hardcoded to ignore any bounding box filled with -1.0.

def create_dataset(image_paths, boxes, classes):
    box_tensor = tf.ragged.constant(boxes, dtype=tf.float32)
    class_tensor = tf.ragged.constant(classes, dtype=tf.float32)
    path_tensor = tf.constant(image_paths)

    dataset = tf.data.Dataset.from_tensor_slices((path_tensor, class_tensor, box_tensor))

    def load_image_and_format(path, cls, box):
        img = tf.io.read_file(path)
        img = tf.image.decode_jpeg(img, channels=3)
        img = tf.cast(img, tf.float32)
        return {"images": img, "bounding_boxes": {"classes": cls, "boxes": box}}

    return dataset.map(load_image_and_format, num_parallel_calls=tf.data.AUTOTUNE)

train_ds = create_dataset(train_img_paths, train_boxes, train_classes)
valid_ds = create_dataset(valid_img_paths, valid_boxes, valid_classes)

# Edge target resolution
resizing_layer = keras_cv.layers.JitteredResize(
    target_size=(416, 416), 
    scale_factor=(0.8, 1.2), 
    bounding_box_format="xyxy" 
)

def prepare_dataset(dataset, batch_size, is_training=True):
    dataset = dataset.ragged_batch(batch_size, drop_remainder=is_training)
    dataset = dataset.map(resizing_layer, num_parallel_calls=tf.data.AUTOTUNE)
    
    # THE XLA FIX: Convert Ragged to Dense by padding with -1.0
    def dict_to_dense_tuple(inputs):
        dense_boxes = inputs["bounding_boxes"]["boxes"].to_tensor(default_value=-1.0)
        dense_classes = inputs["bounding_boxes"]["classes"].to_tensor(default_value=-1.0)
        return inputs["images"], {"boxes": dense_boxes, "classes": dense_classes}
        
    dataset = dataset.map(dict_to_dense_tuple, num_parallel_calls=tf.data.AUTOTUNE)
    return dataset.prefetch(tf.data.AUTOTUNE)

train_dataset = prepare_dataset(train_ds, BATCH_SIZE, is_training=True)
valid_dataset = prepare_dataset(valid_ds, BATCH_SIZE, is_training=False)

Step 3: Building & Training the YOLOv8-XS Model

For edge devices, you want a model with a low parameter count. The yolo_v8_xs_backbone (Extra Small) weighs in at just around 3.4 million parameters.

Pro-Tip: Mobile NPUs panic when they see dynamic tensor shapes. To guarantee compatibility, we pass a dummy_input of pure zeros through the model right after building it. This forces Keras to trace and lock a strictly static computational graph.

import keras

INPUT_SHAPE = (416, 416, 3)

backbone = keras_cv.models.YOLOV8Backbone.from_preset("yolo_v8_xs_backbone")

model = keras_cv.models.YOLOV8Detector(
    num_classes=2,
    bounding_box_format="xyxy",
    backbone=backbone,
    fpn_depth=1, 
)

# Lock in the static shape for Edge compatibility
dummy_input = keras.ops.zeros((1, *INPUT_SHAPE))
_ = model(dummy_input)

# YOLOv8 mathematically requires binary crossentropy for its classification head
optimizer = keras.optimizers.Adam(learning_rate=1e-4)

model.compile(
    optimizer=optimizer,
    classification_loss="binary_crossentropy", 
    box_loss="ciou"  
)

callbacks = [
    keras.callbacks.ModelCheckpoint(
        filepath="best_edge_detector.keras", 
        monitor="val_loss",
        save_best_only=True,
        mode="min",
        verbose=1
    ),
    keras.callbacks.EarlyStopping(
        monitor="val_loss",
        patience=20, # Give it time to learn!
        restore_best_weights=True,
        verbose=1
    )
]

history = model.fit(
    train_dataset,
    validation_data=valid_dataset,
    epochs=100,
    callbacks=callbacks
)

Step 4: Standalone Inference

Once trained, your model is a single, beautiful .keras file. Here is how you run a forward pass using standard OpenCV. Notice the trick to mathematically scale the small 416x416 predicted coordinates back up to fit your original high-resolution image perfectly.

import cv2
import numpy as np
import matplotlib.pyplot as plt

model = keras.models.load_model("best_edge_detector.keras", compile=False) 

def run_inference(image_path):
    img = cv2.imread(image_path)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    orig_h, orig_w = img.shape[:2]
    
    img_resized = cv2.resize(img, (416, 416))
    img_tensor = tf.expand_dims(tf.cast(img_resized, tf.float32), axis=0)
    
    predictions = model.predict(img_tensor, verbose=0)
    
    boxes = predictions['boxes'][0]
    classes = predictions['classes'][0]
    confidences = predictions['confidence'][0]
    
    scale_x = orig_w / 416.0
    scale_y = orig_h / 416.0
    draw_img = img.copy()

    for i in range(len(boxes)):
        if int(classes[i]) == -1 or confidences[i] < 0.5: continue
            
        xmin, ymin, xmax, ymax = boxes[i]
        xmin, ymin = int(xmin * scale_x), int(ymin * scale_y)
        xmax, ymax = int(xmax * scale_x), int(ymax * scale_y)
        
        cv2.rectangle(draw_img, (xmin, ymin), (xmax, ymax), (0, 255, 0), 3)

    plt.imshow(draw_img)
    plt.show()

run_inference("test_image.jpg")

Step 5: The Edge Conversions

Finally, the moment of truth. We need to shrink this model and format it for Apple and Android hardware.

Converting to TFLite (Android / Edge TPU)

We apply INT8 Quantization, converting heavy floating-point math into lightweight 8-bit integers, slashing the model size by ~4x. We also use a brilliant fail-safe: SELECT_TF_OPS. This ensures that if the converter encounters a complex bounding box decoding logic that lacks a strict integer equivalent, it gracefully falls back to standard TensorFlow operations instead of crashing.

converter = tf.lite.TFLiteConverter.from_keras_model(model)

def representative_data_gen():
    image_files = [f for f in os.listdir("Mask-Wearing-19/valid") if f.endswith('.jpg')][:100]
    for img_name in image_files:
        img = cv2.imread(os.path.join("Mask-Wearing-19/valid", img_name))
        img_resized = cv2.resize(cv2.cvtColor(img, cv2.COLOR_BGR2RGB), (416, 416))
        yield [tf.expand_dims(tf.cast(img_resized, tf.float32), axis=0)]

converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.representative_dataset = representative_data_gen
converter.target_spec.supported_ops = [
    tf.lite.OpsSet.TFLITE_BUILTINS_INT8, 
    tf.lite.OpsSet.SELECT_TF_OPS
]
converter.inference_input_type = tf.float32
converter.inference_output_type = tf.float32

tflite_model = converter.convert()
with open("edge_detector_quantized.tflite", "wb") as f:
    f.write(tflite_model)

Converting to CoreML (iOS / Apple Neural Engine)

Apple’s coremltools is notoriously strict about tensor naming. Instead of guessing if Keras named our input layer "inputs" or "keras_tensor_1", we programmatically pull the exact input name directly from the model graph to guarantee a flawless conversion into a modern .mlpackage file.

import coremltools as ct
import shutil
from google.colab import files

# Dynamically grab the entry point name to prevent CoreML crashes
model_input_name = model.inputs[0].name.split(':')[0]

image_input = ct.ImageType(
    name=model_input_name, 
    shape=(1, 416, 416, 3),
    color_layout=ct.colorlayout.RGB,
)

mlmodel = ct.convert(
    model,
    inputs=[image_input],
    convert_to="mlprogram",
    compute_units=ct.ComputeUnit.ALL # Routes math to the Apple Neural Engine
)

mlmodel.save("EdgeDetector.mlpackage")

# Zip it up for easy downloading from Colab/Jupyter
shutil.make_archive("EdgeDetector.mlpackage", 'zip', "EdgeDetector.mlpackage")
files.download("EdgeDetector.mlpackage.zip")

Wrapping Up

And just like that, you have a complete, mathematically sound, zero-bloat pipeline. You’ve navigated XLA ragged tensor compilation, stabilized a YOLO graph for static execution, bypassed restrictive commercial licensing, and successfully converted your weights to both TFLite (havin size of 3.7MB) and CoreML (with size of 6.75 MB with .mlmodel in size of 232kb).

Your model is ready. Now get it into Xcode or Android Studio!