Object detection with Model Garden

View on TensorFlow.org Run in Google Colab View on GitHub Download notebook

This tutorial fine-tunes a RetinaNet with ResNet-50 as backbone model from the TensorFlow Model Garden package (tensorflow-models) to detect three different Blood Cells in BCCD dataset. The RetinaNet is pretrained on COCO train2017 and evaluated on COCO val2017

Model Garden contains a collection of state-of-the-art models, implemented with TensorFlow's high-level APIs. The implementations demonstrate the best practices for modeling, letting users to take full advantage of TensorFlow for their research and product development.

This tutorial demonstrates how to:

  1. Use models from the Tensorflow Model Garden(TFM) package.
  2. Fine-tune a pre-trained RetinanNet with ResNet-50 as backbone for object detection.
  3. Export the tuned RetinaNet model

Install necessary dependencies

pip install -U -q "tf-models-official"

Import required libraries

import os
import io
import pprint
import tempfile
import matplotlib
import numpy as np
import tensorflow as tf
import matplotlib.pyplot as plt

from PIL import Image
from six import BytesIO
from IPython import display
from urllib.request import urlopen

Import required libraries from tensorflow models

import orbit
import tensorflow_models as tfm

from official.core import exp_factory
from official.core import config_definitions as cfg
from official.vision.serving import export_saved_model_lib
from official.vision.ops.preprocess_ops import normalize_image
from official.vision.ops.preprocess_ops import resize_and_crop_image
from official.vision.utils.object_detection import visualization_utils
from official.vision.dataloaders.tf_example_decoder import TfExampleDecoder

pp = pprint.PrettyPrinter(indent=4) # Set Pretty Print Indentation
print(tf.__version__) # Check the version of tensorflow used

%matplotlib inline
/tmpfs/src/tf_docs_env/lib/python3.11/site-packages/tensorflow_hub/__init__.py:61: UserWarning: pkg_resources is deprecated as an API. See https://setuptools.pypa.io/en/latest/pkg_resources.html. The pkg_resources package is slated for removal as early as 2025-11-30. Refrain from using this package or pin to Setuptools<81.
  from pkg_resources import parse_version
2.20.0

Custom dataset preparation for object detection

Models in official repository(of model-garden) requires data in a TFRecords format.

Please check this resource to learn more about TFRecords data format.

Upload your custom data in drive or local disk of the notebook and unzip the data

curl -L 'https://public.roboflow.com/ds/ZpYLqHeT0W?key=ZXfZLRnhsc' > './BCCD.v1-bccd.coco.zip'
unzip -q -o './BCCD.v1-bccd.coco.zip' -d './BCC.v1-bccd.coco/'
rm './BCCD.v1-bccd.coco.zip'
% Total    % Received % Xferd  Average Speed   Time    Time     Time  Current
                                 Dload  Upload   Total   Spent    Left  Speed

  0     0    0     0    0     0      0      0 --:--:-- --:--:-- --:--:--     0
100   901  100   901    0     0   3478      0 --:--:-- --:--:-- --:--:--  3478
85 15.2M   85 13.0M    0     0  23.8M      0 --:--:-- --:--:-- --:--:-- 23.8M
100 15.2M  100 15.2M    0     0  27.4M      0 --:--:-- --:--:-- --:--:--  313M

CLI command to convert data(train data).


TRAIN_DATA_DIR='./BCC.v1-bccd.coco/train'
TRAIN_ANNOTATION_FILE_DIR='./BCC.v1-bccd.coco/train/_annotations.coco.json'
OUTPUT_TFRECORD_TRAIN='./bccd_coco_tfrecords/train'

# Need to provide
  # 1. image_dir: where images are present
  # 2. object_annotations_file: where annotations are listed in json format
  # 3. output_file_prefix: where to write output convered TFRecords files
python -m official.vision.data.create_coco_tf_record --logtostderr \
  --image_dir=${TRAIN_DATA_DIR} \
  --object_annotations_file=${TRAIN_ANNOTATION_FILE_DIR} \
  --output_file_prefix=$OUTPUT_TFRECORD_TRAIN \
  --num_shards=1
Error processing line 1 of /tmpfs/src/tf_docs_env/lib/python3.11/site-packages/tfdocs_sitecustomize.
pth:
Traceback (most recent call last):
    File "<frozen site>", line 195, in addpackage
    File "<st
ring>", line 1
      import _tfdocs_sitecustomize\n
                                   ^
  SyntaxErr
or: unexpected character after line continuation character

Remainder of file ignored
Error processing line 1 of /tmpfs/src/tf_docs_env/lib/python3.11/site-packages/tfdocs_sitecustomize.
pth:

  Traceback (most recent call last):
    File "<frozen site>", line 195, in addpackage
    Fil
e "<string>", line 1
      import _tfdocs_sitecustomize\n
                                   ^
  Syn
taxError: unexpected character after line continuation character

Remainder of file ignored
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
W0000 00:00:1
793882159.877960   32894 gpu_device.cc:2342] Cannot dlopen some GPU libraries. Please make sure the
missing libraries mentioned above are installed properly if you would like to use GPU. Follow the gu
ide at https://www.tensorflow.org/install/gpu for how to download and setup the required libraries f
or your platform.
Skipping registering GPU devices...
I1105 12:35:59.878583 140474347153216 create_coco_tf_record.py:502] writing to output path: ./bccd_c
oco_tfrecords/train
I1105 12:36:00.103268 140474347153216 create_coco_tf_record.py:374] Building bounding box index.
I1105 12:36:00.105362 140474347153216 create_coco_tf_record.py:385] 0 images are missing bboxes.
I1105 12:36:00.366024 140474347153216 tfrecord_lib.py:168] On image 0
I1105 12:36:00.370388 140474347153216 tfrecord_lib.py:168] On image 100
I1105 12:36:00.374054 140474347153216 tfrecord_lib.py:168] On image 200
I1105 12:36:00.377497 140474347153216 tfrecord_lib.py:168] On image 300
I1105 12:36:00.380959 140474347153216 tfrecord_lib.py:168] On image 400
I1105 12:36:00.384281 140474347153216 tfrecord_lib.py:168] On image 500
I1105 12:36:00.387392 140474347153216 tfrecord_lib.py:168] On image 600
I1105 12:36:00.390604 140474347153216 tfrecord_lib.py:168] On image 700
I1105 12:36:00.417978 140474347153216 tfrecord_lib.py:180] Finished writing, skipped 6 annotations.
I1105 12:36:00.419922 140474347153216 create_coco_tf_record.py:537] Finished writing, skipped 6 anno
tations.

CLI command to convert data(validation data).


VALID_DATA_DIR='./BCC.v1-bccd.coco/valid'
VALID_ANNOTATION_FILE_DIR='./BCC.v1-bccd.coco/valid/_annotations.coco.json'
OUTPUT_TFRECORD_VALID='./bccd_coco_tfrecords/valid'

python -m official.vision.data.create_coco_tf_record --logtostderr \
  --image_dir=$VALID_DATA_DIR \
  --object_annotations_file=$VALID_ANNOTATION_FILE_DIR \
  --output_file_prefix=$OUTPUT_TFRECORD_VALID \
  --num_shards=1
Error processing line 1 of /tmpfs/src/tf_docs_env/lib/python3.11/site-packages/tfdocs_sitecustomize.
pth:
Traceback (most recent call last):
    File "<frozen site>", line 195, in addpackage
    File "<st
ring>", line 1
      import _tfdocs_sitecustomize\n
                                   ^
  SyntaxErr
or: unexpected character after line continuation character

Remainder of file ignored
Error processing line 1 of /tmpfs/src/tf_docs_env/lib/python3.11/site-packages/tfdocs_sitecustomize.
pth:

  Traceback (most recent call last):
    File "<frozen site>", line 195, in addpackage
    Fil
e "<string>", line 1
      import _tfdocs_sitecustomize\n
                                   ^
  Syn
taxError: unexpected character after line continuation character

Remainder of file ignored
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
W0000 00:00:1
793882166.487063   33040 gpu_device.cc:2342] Cannot dlopen some GPU libraries. Please make sure the
missing libraries mentioned above are installed properly if you would like to use GPU. Follow the gu
ide at https://www.tensorflow.org/install/gpu for how to download and setup the required libraries f
or your platform.
Skipping registering GPU devices...
I1105 12:36:06.487452 140335292204864 create_coco_tf_record.py:502] writing to output path: ./bccd_c
oco_tfrecords/valid
I1105 12:36:06.497694 140335292204864 create_coco_tf_record.py:374] Building bounding box index.
I1105 12:36:06.497938 140335292204864 create_coco_tf_record.py:385] 0 images are missing bboxes.
I1105 12:36:06.711694 140335292204864 tfrecord_lib.py:168] On image 0
I1105 12:36:06.739617 140335292204864 tfrecord_lib.py:180] Finished writing, skipped 0 annotations.
I1105 12:36:06.740615 140335292204864 create_coco_tf_record.py:537] Finished writing, skipped 0 anno
tations.

CLI command to convert data(test data).


TEST_DATA_DIR='./BCC.v1-bccd.coco/test'
TEST_ANNOTATION_FILE_DIR='./BCC.v1-bccd.coco/test/_annotations.coco.json'
OUTPUT_TFRECORD_TEST='./bccd_coco_tfrecords/test'

python -m official.vision.data.create_coco_tf_record --logtostderr \
  --image_dir=$TEST_DATA_DIR \
  --object_annotations_file=$TEST_ANNOTATION_FILE_DIR \
  --output_file_prefix=$OUTPUT_TFRECORD_TEST \
  --num_shards=1
Error processing line 1 of /tmpfs/src/tf_docs_env/lib/python3.11/site-packages/tfdocs_sitecustomize.
pth:
Traceback (most recent call last):
    File "<frozen site>", line 195, in addpackage
    File "<st
ring>", line 1
      import _tfdocs_sitecustomize\n
                                   ^
  SyntaxErr
or: unexpected character after line continuation character

Remainder of file ignored
Error processing line 1 of /tmpfs/src/tf_docs_env/lib/python3.11/site-packages/tfdocs_sitecustomize.
pth:

  Traceback (most recent call last):
    File "<frozen site>", line 195, in addpackage
    Fil
e "<string>", line 1
      import _tfdocs_sitecustomize\n
                                   ^
  Syn
taxError: unexpected character after line continuation character

Remainder of file ignored
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
W0000 00:00:1
793882172.794489   33186 gpu_device.cc:2342] Cannot dlopen some GPU libraries. Please make sure the
missing libraries mentioned above are installed properly if you would like to use GPU. Follow the gu
ide at https://www.tensorflow.org/install/gpu for how to download and setup the required libraries f
or your platform.
Skipping registering GPU devices...
I1105 12:36:12.794876 140423524677440 create_coco_tf_record.py:502] writing to output path: ./bccd_c
oco_tfrecords/test
I1105 12:36:12.800759 140423524677440 create_coco_tf_record.py:374] Building bounding box index.
I1105 12:36:12.801036 140423524677440 create_coco_tf_record.py:385] 0 images are missing bboxes.
I1105 12:36:13.006542 140423524677440 tfrecord_lib.py:168] On image 0
I1105 12:36:13.032616 140423524677440 tfrecord_lib.py:180] Finished writing, skipped 0 annotations.
I1105 12:36:13.033573 140423524677440 create_coco_tf_record.py:537] Finished writing, skipped 0 anno
tations.

Configure the Retinanet Resnet FPN COCO model for custom dataset.

Dataset used for fine tuning the checkpoint is Blood Cells Detection (BCCD).

train_data_input_path = './bccd_coco_tfrecords/train-00000-of-00001.tfrecord'
valid_data_input_path = './bccd_coco_tfrecords/valid-00000-of-00001.tfrecord'
test_data_input_path = './bccd_coco_tfrecords/test-00000-of-00001.tfrecord'
model_dir = './trained_model/'
export_dir ='./exported_model/'

In Model Garden, the collections of parameters that define a model are called configs. Model Garden can create a config based on a known set of parameters via a factory.

Use the retinanet_resnetfpn_coco experiment configuration, as defined by tfm.vision.configs.retinanet.retinanet_resnetfpn_coco.

The configuration defines an experiment to train a RetinanNet with Resnet-50 as backbone, FPN as decoder. Default Configuration is trained on COCO train2017 and evaluated on COCO val2017.

There are also other alternative experiments available such as retinanet_resnetfpn_coco, retinanet_spinenet_coco, fasterrcnn_resnetfpn_coco and more. One can switch to them by changing the experiment name argument to the get_exp_config function.

We are going to fine tune the Resnet-50 backbone checkpoint which is already present in the default configuration.

exp_config = exp_factory.get_exp_config('retinanet_resnetfpn_coco')

Adjust the model and dataset configurations so that it works with custom dataset(in this case BCCD).

batch_size = 8
num_classes = 3

HEIGHT, WIDTH = 256, 256
IMG_SIZE = [HEIGHT, WIDTH, 3]

# Backbone config.
exp_config.task.freeze_backbone = False
exp_config.task.annotation_file = ''

# Model config.
exp_config.task.model.input_size = IMG_SIZE
exp_config.task.model.num_classes = num_classes + 1
exp_config.task.model.detection_generator.tflite_post_processing.max_classes_per_detection = exp_config.task.model.num_classes

# Training data config.
exp_config.task.train_data.input_path = train_data_input_path
exp_config.task.train_data.dtype = 'float32'
exp_config.task.train_data.global_batch_size = batch_size
exp_config.task.train_data.parser.aug_scale_max = 1.0
exp_config.task.train_data.parser.aug_scale_min = 1.0

# Validation data config.
exp_config.task.validation_data.input_path = valid_data_input_path
exp_config.task.validation_data.dtype = 'float32'
exp_config.task.validation_data.global_batch_size = batch_size

Adjust the trainer configuration.

logical_device_names = [logical_device.name for logical_device in tf.config.list_logical_devices()]

if 'GPU' in ''.join(logical_device_names):
  print('This may be broken in Colab.')
  device = 'GPU'
elif 'TPU' in ''.join(logical_device_names):
  print('This may be broken in Colab.')
  device = 'TPU'
else:
  print('Running on CPU is slow, so only train for a few steps.')
  device = 'CPU'


train_steps = 1000
exp_config.trainer.steps_per_loop = 100 # steps_per_loop = num_of_training_examples // train_batch_size

exp_config.trainer.summary_interval = 100
exp_config.trainer.checkpoint_interval = 100
exp_config.trainer.validation_interval = 100
exp_config.trainer.validation_steps =  100 # validation_steps = num_of_validation_examples // eval_batch_size
exp_config.trainer.train_steps = train_steps
exp_config.trainer.optimizer_config.warmup.linear.warmup_steps = 100
exp_config.trainer.optimizer_config.learning_rate.type = 'cosine'
exp_config.trainer.optimizer_config.learning_rate.cosine.decay_steps = train_steps
exp_config.trainer.optimizer_config.learning_rate.cosine.initial_learning_rate = 0.1
exp_config.trainer.optimizer_config.warmup.linear.warmup_learning_rate = 0.05
Running on CPU is slow, so only train for a few steps.
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
W0000 00:00:1793882174.538177   32689 gpu_device.cc:2342] Cannot dlopen some GPU libraries. Please make sure the missing libraries mentioned above are installed properly if you would like to use GPU. Follow the guide at https://www.tensorflow.org/install/gpu for how to download and setup the required libraries for your platform.
Skipping registering GPU devices...
pp.pprint(exp_config.as_dict())
display.Javascript('google.colab.output.setIframeHeight("500px");')
{   'runtime': {   'all_reduce_alg': None,
                   'batchnorm_spatial_persistent': False,
                   'dataset_num_private_threads': None,
                   'default_shard_dim': -1,
                   'distribution_strategy': 'mirrored',
                   'enable_xla': False,
                   'gpu_thread_mode': None,
                   'loss_scale': None,
                   'mixed_precision_dtype': 'bfloat16',
                   'num_cores_per_replica': 1,
                   'num_gpus': 0,
                   'num_packs': 1,
                   'per_gpu_thread_count': 0,
                   'run_eagerly': False,
                   'task_index': -1,
                   'tpu': None,
                   'tpu_enable_xla_dynamic_padder': None,
                   'use_tpu_mp_strategy': False,
                   'worker_hosts': None},
    'task': {   'allow_image_summary': False,
                'annotation_file': '',
                'differential_privacy_config': None,
                'export_config': {   'cast_detection_classes_to_float': False,
                                     'cast_num_detections_to_float': False,
                                     'output_intermediate_features': False,
                                     'output_normalized_coordinates': False},
                'freeze_backbone': False,
                'init_checkpoint': 'gs://cloud-tpu-checkpoints/vision-2.0/resnet50_imagenet/ckpt-28080',
                'init_checkpoint_modules': 'backbone',
                'losses': {   'box_loss_weight': 50,
                              'focal_loss_alpha': 0.25,
                              'focal_loss_gamma': 1.5,
                              'huber_loss_delta': 0.1,
                              'l2_weight_decay': 0.0001,
                              'loss_weight': 1.0},
                'max_num_eval_detections': 100,
                'model': {   'anchor': {   'anchor_size': 4.0,
                                           'aspect_ratios': [0.5, 1.0, 2.0],
                                           'num_scales': 3},
                             'backbone': {   'resnet': {   'bn_trainable': True,
                                                           'depth_multiplier': 1.0,
                                                           'model_id': 50,
                                                           'replace_stem_max_pool': False,
                                                           'resnetd_shortcut': False,
                                                           'scale_stem': True,
                                                           'se_ratio': 0.0,
                                                           'stem_type': 'v0',
                                                           'stochastic_depth_drop_rate': 0.0},
                                             'type': 'resnet'},
                             'decoder': {   'fpn': {   'fusion_type': 'sum',
                                                       'num_filters': 256,
                                                       'use_keras_layer': False,
                                                       'use_separable_conv': False},
                                            'type': 'fpn'},
                             'detection_generator': {   'apply_nms': True,
                                                        'box_coder_weights': None,
                                                        'decode_boxes': True,
                                                        'max_num_detections': 100,
                                                        'nms_iou_threshold': 0.5,
                                                        'nms_version': 'v2',
                                                        'pre_nms_score_threshold': 0.05,
                                                        'pre_nms_top_k': 5000,
                                                        'return_decoded': None,
                                                        'soft_nms_sigma': None,
                                                        'tflite_post_processing': {   'detections_per_class': 5,
                                                                                      'h_scale': 1.0,
                                                                                      'max_classes_per_detection': 4,
                                                                                      'max_detections': 200,
                                                                                      'nms_iou_threshold': 0.5,
                                                                                      'nms_score_threshold': 0.1,
                                                                                      'normalize_anchor_coordinates': False,
                                                                                      'omit_nms': False,
                                                                                      'use_regular_nms': False,
                                                                                      'w_scale': 1.0,
                                                                                      'x_scale': 1.0,
                                                                                      'y_scale': 1.0},
                                                        'use_class_agnostic_nms': False,
                                                        'use_cpu_nms': False},
                             'head': {   'attribute_heads': [],
                                         'num_convs': 4,
                                         'num_filters': 256,
                                         'share_classification_heads': False,
                                         'share_level_convs': True,
                                         'use_separable_conv': False},
                             'input_size': [256, 256, 3],
                             'max_level': 7,
                             'min_level': 3,
                             'norm_activation': {   'activation': 'relu',
                                                    'norm_epsilon': 0.001,
                                                    'norm_momentum': 0.99,
                                                    'use_sync_bn': False},
                             'num_classes': 4},
                'name': None,
                'per_category_metrics': False,
                'train_data': {   'apply_tf_data_service_before_batching': False,
                                  'autotune_algorithm': None,
                                  'block_length': 1,
                                  'cache': False,
                                  'cycle_length': None,
                                  'decoder': {   'simple_decoder': {   'attribute_names': [   ],
                                                                       'mask_binarize_threshold': None,
                                                                       'regenerate_source_id': False},
                                                 'type': 'simple_decoder'},
                                  'deterministic': None,
                                  'drop_remainder': True,
                                  'dtype': 'float32',
                                  'enable_shared_tf_data_service_between_parallel_trainers': False,
                                  'enable_tf_data_service': False,
                                  'file_type': 'tfrecord',
                                  'global_batch_size': 8,
                                  'input_path': './bccd_coco_tfrecords/train-00000-of-00001.tfrecord',
                                  'is_training': True,
                                  'parser': {   'aug_policy': None,
                                                'aug_rand_hflip': True,
                                                'aug_rand_jpeg': None,
                                                'aug_scale_max': 1.0,
                                                'aug_scale_min': 1.0,
                                                'aug_type': None,
                                                'keep_aspect_ratio': True,
                                                'match_threshold': 0.5,
                                                'max_num_instances': 100,
                                                'num_channels': 3,
                                                'pad': True,
                                                'skip_crowd_during_training': True,
                                                'unmatched_threshold': 0.5},
                                  'prefetch_buffer_size': None,
                                  'ram_budget': None,
                                  'seed': None,
                                  'sharding': True,
                                  'shuffle_buffer_size': 10000,
                                  'tf_data_service_address': None,
                                  'tf_data_service_job_name': None,
                                  'tfds_as_supervised': False,
                                  'tfds_data_dir': '',
                                  'tfds_name': '',
                                  'tfds_skip_decoding_feature': '',
                                  'tfds_split': '',
                                  'trainer_id': None,
                                  'weights': None},
                'use_coco_metrics': True,
                'use_wod_metrics': False,
                'validation_data': {   'apply_tf_data_service_before_batching': False,
                                       'autotune_algorithm': None,
                                       'block_length': 1,
                                       'cache': False,
                                       'cycle_length': None,
                                       'decoder': {   'simple_decoder': {   'attribute_names': [   ],
                                                                            'mask_binarize_threshold': None,
                                                                            'regenerate_source_id': False},
                                                      'type': 'simple_decoder'},
                                       'deterministic': None,
                                       'drop_remainder': True,
                                       'dtype': 'float32',
                                       'enable_shared_tf_data_service_between_parallel_trainers': False,
                                       'enable_tf_data_service': False,
                                       'file_type': 'tfrecord',
                                       'global_batch_size': 8,
                                       'input_path': './bccd_coco_tfrecords/valid-00000-of-00001.tfrecord',
                                       'is_training': False,
                                       'parser': {   'aug_policy': None,
                                                     'aug_rand_hflip': False,
                                                     'aug_rand_jpeg': None,
                                                     'aug_scale_max': 1.0,
                                                     'aug_scale_min': 1.0,
                                                     'aug_type': None,
                                                     'keep_aspect_ratio': True,
                                                     'match_threshold': 0.5,
                                                     'max_num_instances': 100,
                                                     'num_channels': 3,
                                                     'pad': True,
                                                     'skip_crowd_during_training': True,
                                                     'unmatched_threshold': 0.5},
                                       'prefetch_buffer_size': None,
                                       'ram_budget': None,
                                       'seed': None,
                                       'sharding': True,
                                       'shuffle_buffer_size': 10000,
                                       'tf_data_service_address': None,
                                       'tf_data_service_job_name': None,
                                       'tfds_as_supervised': False,
                                       'tfds_data_dir': '',
                                       'tfds_name': '',
                                       'tfds_skip_decoding_feature': '',
                                       'tfds_split': '',
                                       'trainer_id': None,
                                       'weights': None} },
    'trainer': {   'allow_tpu_summary': False,
                   'best_checkpoint_eval_metric': '',
                   'best_checkpoint_export_subdir': '',
                   'best_checkpoint_metric_comp': 'higher',
                   'checkpoint_interval': 100,
                   'continuous_eval_timeout': 3600,
                   'eval_tf_function': True,
                   'eval_tf_while_loop': False,
                   'loss_upper_bound': 1000000.0,
                   'max_to_keep': 5,
                   'optimizer_config': {   'ema': None,
                                           'learning_rate': {   'cosine': {   'alpha': 0.0,
                                                                              'decay_steps': 1000,
                                                                              'initial_learning_rate': 0.1,
                                                                              'name': 'CosineDecay',
                                                                              'offset': 0},
                                                                'type': 'cosine'},
                                           'optimizer': {   'sgd': {   'clipnorm': None,
                                                                       'clipvalue': None,
                                                                       'decay': 0.0,
                                                                       'global_clipnorm': None,
                                                                       'momentum': 0.9,
                                                                       'name': 'SGD',
                                                                       'nesterov': False},
                                                            'type': 'sgd'},
                                           'warmup': {   'linear': {   'name': 'linear',
                                                                       'warmup_learning_rate': 0.05,
                                                                       'warmup_steps': 100},
                                                         'type': 'linear'} },
                   'preemption_on_demand_checkpoint': True,
                   'recovery_begin_steps': 0,
                   'recovery_max_trials': 0,
                   'steps_per_loop': 100,
                   'summary_interval': 100,
                   'train_steps': 1000,
                   'train_tf_function': True,
                   'train_tf_while_loop': True,
                   'validation_interval': 100,
                   'validation_steps': 100,
                   'validation_summary_subdir': 'validation'} }
<IPython.core.display.Javascript object>

Set up the distribution strategy.

if exp_config.runtime.mixed_precision_dtype == tf.float16:
    tf.keras.mixed_precision.set_global_policy('mixed_float16')

if 'GPU' in ''.join(logical_device_names):
  distribution_strategy = tf.distribute.MirroredStrategy()
elif 'TPU' in ''.join(logical_device_names):
  tf.tpu.experimental.initialize_tpu_system()
  tpu = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='/device:TPU_SYSTEM:0')
  distribution_strategy = tf.distribute.experimental.TPUStrategy(tpu)
else:
  print('Warning: this will be really slow.')
  distribution_strategy = tf.distribute.OneDeviceStrategy(logical_device_names[0])

print('Done')
Warning: this will be really slow.
Done

Create the Task object (tfm.core.base_task.Task) from the config_definitions.TaskConfig.

The Task object has all the methods necessary for building the dataset, building the model, and running training & evaluation. These methods are driven by tfm.core.train_lib.run_experiment.

with distribution_strategy.scope():
  task = tfm.core.task_factory.get_task(exp_config.task, logging_dir=model_dir)

Visualize a batch of the data.

for images, labels in task.build_inputs(exp_config.task.train_data).take(1):
  print()
  print(f'images.shape: {str(images.shape):16}  images.dtype: {images.dtype!r}')
  print(f'labels.keys: {labels.keys()}')
images.shape: (8, 256, 256, 3)  images.dtype: tf.float32
labels.keys: dict_keys(['cls_targets', 'box_targets', 'anchor_boxes', 'cls_weights', 'box_weights', 'image_info'])

Create category index dictionary to map the labels to coressponding label names.

category_index={
    1: {
        'id': 1,
        'name': 'Platelets'
       },
    2: {
        'id': 2,
        'name': 'RBC'
       },
    3: {
        'id': 3,
        'name': 'WBC'
       }
}
tf_ex_decoder = TfExampleDecoder()

Helper function for visualizing the results from TFRecords.

Use visualize_boxes_and_labels_on_image_array from visualization_utils to draw boudning boxes on the image.

def show_batch(raw_records, num_of_examples):
  plt.figure(figsize=(20, 20))
  use_normalized_coordinates=True
  min_score_thresh = 0.30
  for i, serialized_example in enumerate(raw_records):
    plt.subplot(1, 3, i + 1)
    decoded_tensors = tf_ex_decoder.decode(serialized_example)
    image = decoded_tensors['image'].numpy().astype('uint8')
    scores = np.ones(shape=(len(decoded_tensors['groundtruth_boxes'])))
    visualization_utils.visualize_boxes_and_labels_on_image_array(
        image,
        decoded_tensors['groundtruth_boxes'].numpy(),
        decoded_tensors['groundtruth_classes'].numpy().astype('int'),
        scores,
        category_index=category_index,
        use_normalized_coordinates=use_normalized_coordinates,
        max_boxes_to_draw=200,
        min_score_thresh=min_score_thresh,
        agnostic_mode=False,
        instance_masks=None,
        line_thickness=4)

    plt.imshow(image)
    plt.axis('off')
    plt.title(f'Image-{i+1}')
  plt.show()

Visualization of train data

The bounding box detection has two components

  1. Class label of the object detected (e.g.RBC)
  2. Percentage of match between predicted and ground truth bounding boxes.
buffer_size = 20
num_of_examples = 3

raw_records = tf.data.TFRecordDataset(
    exp_config.task.train_data.input_path).shuffle(
        buffer_size=buffer_size).take(num_of_examples)
show_batch(raw_records, num_of_examples)

png

Train and evaluate.

We follow the COCO challenge tradition to evaluate the accuracy of object detection based on mAP(mean Average Precision). Please check here for detail explanation of how evaluation metrics for detection task is done.

IoU: is defined as the area of the intersection divided by the area of the union of a predicted bounding box and ground truth bounding box.

model, eval_logs = tfm.core.train_lib.run_experiment(
    distribution_strategy=distribution_strategy,
    task=task,
    mode='train_and_eval',
    params=exp_config,
    model_dir=model_dir,
    run_post_eval=True)
restoring or initializing model...
INFO:tensorflow:Customized initialization is done through the passed `init_fn`.
INFO:tensorflow:Customized initialization is done through the passed `init_fn`.
train | step:      0 | training until step 100...
2026-11-05 12:36:30.378476: W tensorflow/core/framework/dataset.cc:993] Input of GeneratorDatasetOp::Dataset will not be optimized because the dataset does not implement the AsGraphDefInternal() method needed to apply optimizations.
train | step:    100 | steps/sec:    0.7 | output: 
    {'box_loss': np.float32(0.02235822),
     'cls_loss': np.float32(0.7599054),
     'learning_rate': np.float32(0.09755283),
     'model_loss': np.float32(1.8778163),
     'total_loss': np.float32(2.7424054),
     'training_loss': np.float32(2.7424054)}
saved checkpoint to ./trained_model/ckpt-100.
 eval | step:    100 | running 100 steps of evaluation...
2026-11-05 12:38:50.432662: W tensorflow/core/framework/dataset.cc:993] Input of GeneratorDatasetOp::Dataset will not be optimized because the dataset does not implement the AsGraphDefInternal() method needed to apply optimizations.
creating index...
index created!
creating index...
index created!
Running per image evaluation...
Evaluate annotation type *bbox*
DONE (t=0.67s).
Accumulating evaluation results...
DONE (t=0.05s).
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.000
 Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets=100 ] = 0.002
 Average Precision  (AP) @[ IoU=0.75      | area=   all | maxDets=100 ] = 0.000
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.000
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.000
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.000
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=  1 ] = 0.000
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets= 10 ] = 0.001
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.009
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.000
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.010
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.000
 eval | step:    100 | steps/sec:    7.6 | eval time:   13.1 sec | output: 
    {'AP': np.float32(0.00028884076),
     'AP50': np.float32(0.0017419293),
     'AP75': np.float32(4.4976234e-05),
     'APl': np.float32(0.0),
     'APm': np.float32(0.00028936844),
     'APs': np.float32(0.0),
     'ARl': np.float32(0.0),
     'ARm': np.float32(0.010283319),
     'ARmax1': np.float32(0.0),
     'ARmax10': np.float32(0.0009075908),
     'ARmax100': np.float32(0.009266979),
     'ARs': np.float32(0.0),
     'box_loss': np.float32(0.067625485),
     'cls_loss': np.float32(2.406881e+06),
     'model_loss': np.float32(2.4068845e+06),
     'steps_per_second': 7.609833056827091,
     'total_loss': np.float32(2.4068852e+06),
     'validation_loss': np.float32(2.4068852e+06)}
train | step:    100 | training until step 200...
train | step:    200 | steps/sec:    0.7 | output: 
    {'box_loss': np.float32(0.008399598),
     'cls_loss': np.float32(0.5154132),
     'learning_rate': np.float32(0.090450846),
     'model_loss': np.float32(0.9353929),
     'total_loss': np.float32(1.9638149),
     'training_loss': np.float32(1.9638149)}
saved checkpoint to ./trained_model/ckpt-200.
 eval | step:    200 | running 100 steps of evaluation...
2026-11-05 12:41:10.207266: W tensorflow/core/framework/dataset.cc:993] Input of GeneratorDatasetOp::Dataset will not be optimized because the dataset does not implement the AsGraphDefInternal() method needed to apply optimizations.
creating index...
index created!
creating index...
index created!
Running per image evaluation...
Evaluate annotation type *bbox*
DONE (t=0.98s).
Accumulating evaluation results...
DONE (t=0.05s).
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.003
 Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets=100 ] = 0.009
 Average Precision  (AP) @[ IoU=0.75      | area=   all | maxDets=100 ] = 0.001
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.000
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.003
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.008
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=  1 ] = 0.015
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets= 10 ] = 0.052
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.124
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.000
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.144
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.130
 eval | step:    200 | steps/sec:   13.0 | eval time:    7.7 sec | output: 
    {'AP': np.float32(0.002787114),
     'AP50': np.float32(0.009249211),
     'AP75': np.float32(0.0006791693),
     'APl': np.float32(0.007855181),
     'APm': np.float32(0.0027444179),
     'APs': np.float32(0.0),
     'ARl': np.float32(0.13023679),
     'ARm': np.float32(0.14424771),
     'ARmax1': np.float32(0.014976409),
     'ARmax10': np.float32(0.0518146),
     'ARmax100': np.float32(0.12352531),
     'ARs': np.float32(0.0),
     'box_loss': np.float32(0.008891074),
     'cls_loss': np.float32(44.435043),
     'model_loss': np.float32(44.879593),
     'steps_per_second': 12.987148021236814,
     'total_loss': np.float32(45.898396),
     'validation_loss': np.float32(45.898396)}
train | step:    200 | training until step 300...
train | step:    300 | steps/sec:    0.8 | output: 
    {'box_loss': np.float32(0.008243298),
     'cls_loss': np.float32(0.48400727),
     'learning_rate': np.float32(0.07938927),
     'model_loss': np.float32(0.89617234),
     'total_loss': np.float32(1.9063846),
     'training_loss': np.float32(1.9063846)}
saved checkpoint to ./trained_model/ckpt-300.
 eval | step:    300 | running 100 steps of evaluation...
2026-11-05 12:43:22.965131: W tensorflow/core/framework/dataset.cc:993] Input of GeneratorDatasetOp::Dataset will not be optimized because the dataset does not implement the AsGraphDefInternal() method needed to apply optimizations.
creating index...
index created!
creating index...
index created!
Running per image evaluation...
Evaluate annotation type *bbox*
DONE (t=1.36s).
Accumulating evaluation results...
DONE (t=0.05s).
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.058
 Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets=100 ] = 0.191
 Average Precision  (AP) @[ IoU=0.75      | area=   all | maxDets=100 ] = 0.004
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.000
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.016
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.099
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=  1 ] = 0.081
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets= 10 ] = 0.100
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.177
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.000
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.112
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.177
 eval | step:    300 | steps/sec:   12.3 | eval time:    8.2 sec | output: 
    {'AP': np.float32(0.058184084),
     'AP50': np.float32(0.19077206),
     'AP75': np.float32(0.004131129),
     'APl': np.float32(0.09912344),
     'APm': np.float32(0.016192215),
     'APs': np.float32(0.0),
     'ARl': np.float32(0.17723133),
     'ARm': np.float32(0.11169576),
     'ARmax1': np.float32(0.081328735),
     'ARmax10': np.float32(0.1003248),
     'ARmax100': np.float32(0.17701623),
     'ARs': np.float32(0.0),
     'box_loss': np.float32(0.007695793),
     'cls_loss': np.float32(0.55965734),
     'model_loss': np.float32(0.944447),
     'steps_per_second': 12.256823500653343,
     'total_loss': np.float32(1.9465537),
     'validation_loss': np.float32(1.9465537)}
train | step:    300 | training until step 400...
train | step:    400 | steps/sec:    0.8 | output: 
    {'box_loss': np.float32(0.007889095),
     'cls_loss': np.float32(0.44259688),
     'learning_rate': np.float32(0.06545085),
     'model_loss': np.float32(0.8370516),
     'total_loss': np.float32(1.8320946),
     'training_loss': np.float32(1.8320946)}
saved checkpoint to ./trained_model/ckpt-400.
 eval | step:    400 | running 100 steps of evaluation...
2026-11-05 12:45:36.397352: W tensorflow/core/framework/dataset.cc:993] Input of GeneratorDatasetOp::Dataset will not be optimized because the dataset does not implement the AsGraphDefInternal() method needed to apply optimizations.
creating index...
index created!
creating index...
index created!
Running per image evaluation...
Evaluate annotation type *bbox*
DONE (t=1.33s).
Accumulating evaluation results...
DONE (t=0.05s).
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.021
 Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets=100 ] = 0.076
 Average Precision  (AP) @[ IoU=0.75      | area=   all | maxDets=100 ] = 0.002
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.000
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.021
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.127
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=  1 ] = 0.004
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets= 10 ] = 0.028
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.114
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.000
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.114
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.128
 eval | step:    400 | steps/sec:   12.7 | eval time:    7.9 sec | output: 
    {'AP': np.float32(0.020811269),
     'AP50': np.float32(0.076047935),
     'AP75': np.float32(0.002485574),
     'APl': np.float32(0.12739274),
     'APm': np.float32(0.020686641),
     'APs': np.float32(0.0),
     'ARl': np.float32(0.12777779),
     'ARm': np.float32(0.11408977),
     'ARmax1': np.float32(0.004166667),
     'ARmax10': np.float32(0.028176568),
     'ARmax100': np.float32(0.11419142),
     'ARs': np.float32(0.0),
     'box_loss': np.float32(0.0076991236),
     'cls_loss': np.float32(0.70400155),
     'model_loss': np.float32(1.0889578),
     'steps_per_second': 12.73849965260189,
     'total_loss': np.float32(2.0771887),
     'validation_loss': np.float32(2.0771887)}
train | step:    400 | training until step 500...
train | step:    500 | steps/sec:    0.8 | output: 
    {'box_loss': np.float32(0.0067797755),
     'cls_loss': np.float32(0.4108571),
     'learning_rate': np.float32(0.049999997),
     'model_loss': np.float32(0.7498457),
     'total_loss': np.float32(1.7323052),
     'training_loss': np.float32(1.7323052)}
saved checkpoint to ./trained_model/ckpt-500.
 eval | step:    500 | running 100 steps of evaluation...
2026-11-05 12:47:48.832966: W tensorflow/core/framework/dataset.cc:993] Input of GeneratorDatasetOp::Dataset will not be optimized because the dataset does not implement the AsGraphDefInternal() method needed to apply optimizations.
creating index...
index created!
creating index...
index created!
Running per image evaluation...
Evaluate annotation type *bbox*
DONE (t=1.31s).
Accumulating evaluation results...
DONE (t=0.05s).
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.111
 Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets=100 ] = 0.305
 Average Precision  (AP) @[ IoU=0.75      | area=   all | maxDets=100 ] = 0.035
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.000
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.032
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.208
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=  1 ] = 0.099
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets= 10 ] = 0.137
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.201
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.000
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.108
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.225
 eval | step:    500 | steps/sec:   12.6 | eval time:    7.9 sec | output: 
    {'AP': np.float32(0.110897474),
     'AP50': np.float32(0.3050338),
     'AP75': np.float32(0.035085954),
     'APl': np.float32(0.20756923),
     'APm': np.float32(0.03220341),
     'APs': np.float32(0.0),
     'ARl': np.float32(0.2248634),
     'ARm': np.float32(0.10793849),
     'ARmax1': np.float32(0.09885709),
     'ARmax10': np.float32(0.13747095),
     'ARmax100': np.float32(0.20096105),
     'ARs': np.float32(0.0),
     'box_loss': np.float32(0.0077647995),
     'cls_loss': np.float32(0.46571958),
     'model_loss': np.float32(0.8539595),
     'steps_per_second': 12.612088047764415,
     'total_loss': np.float32(1.8309981),
     'validation_loss': np.float32(1.8309981)}
train | step:    500 | training until step 600...
train | step:    600 | steps/sec:    0.8 | output: 
    {'box_loss': np.float32(0.0062877936),
     'cls_loss': np.float32(0.38793606),
     'learning_rate': np.float32(0.034549143),
     'model_loss': np.float32(0.7023259),
     'total_loss': np.float32(1.6749988),
     'training_loss': np.float32(1.6749988)}
saved checkpoint to ./trained_model/ckpt-600.
 eval | step:    600 | running 100 steps of evaluation...
2026-11-05 12:50:00.868196: W tensorflow/core/framework/dataset.cc:993] Input of GeneratorDatasetOp::Dataset will not be optimized because the dataset does not implement the AsGraphDefInternal() method needed to apply optimizations.
creating index...
index created!
creating index...
index created!
Running per image evaluation...
Evaluate annotation type *bbox*
DONE (t=1.29s).
Accumulating evaluation results...
DONE (t=0.05s).
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.158
 Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets=100 ] = 0.390
 Average Precision  (AP) @[ IoU=0.75      | area=   all | maxDets=100 ] = 0.069
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.000
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.048
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.257
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=  1 ] = 0.133
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets= 10 ] = 0.184
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.244
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.000
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.119
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.273
 eval | step:    600 | steps/sec:   11.8 | eval time:    8.5 sec | output: 
    {'AP': np.float32(0.15807459),
     'AP50': np.float32(0.3899082),
     'AP75': np.float32(0.06872952),
     'APl': np.float32(0.25653425),
     'APm': np.float32(0.047731973),
     'APs': np.float32(0.0),
     'ARl': np.float32(0.27258652),
     'ARm': np.float32(0.11936825),
     'ARmax1': np.float32(0.13289581),
     'ARmax10': np.float32(0.1844379),
     'ARmax100': np.float32(0.24384384),
     'ARs': np.float32(0.0),
     'box_loss': np.float32(0.0071177282),
     'cls_loss': np.float32(0.43631196),
     'model_loss': np.float32(0.7921984),
     'steps_per_second': 11.776224929921238,
     'total_loss': np.float32(1.7609191),
     'validation_loss': np.float32(1.7609191)}
train | step:    600 | training until step 700...
train | step:    700 | steps/sec:    0.8 | output: 
    {'box_loss': np.float32(0.0058253678),
     'cls_loss': np.float32(0.3767513),
     'learning_rate': np.float32(0.02061074),
     'model_loss': np.float32(0.66801953),
     'total_loss': np.float32(1.633805),
     'training_loss': np.float32(1.633805)}
saved checkpoint to ./trained_model/ckpt-700.
 eval | step:    700 | running 100 steps of evaluation...
2026-11-05 12:52:14.090843: W tensorflow/core/framework/dataset.cc:993] Input of GeneratorDatasetOp::Dataset will not be optimized because the dataset does not implement the AsGraphDefInternal() method needed to apply optimizations.
creating index...
index created!
creating index...
index created!
Running per image evaluation...
Evaluate annotation type *bbox*
DONE (t=1.25s).
Accumulating evaluation results...
DONE (t=0.05s).
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.224
 Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets=100 ] = 0.474
 Average Precision  (AP) @[ IoU=0.75      | area=   all | maxDets=100 ] = 0.174
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.000
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.084
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.318
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=  1 ] = 0.169
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets= 10 ] = 0.246
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.304
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.000
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.150
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.335
 eval | step:    700 | steps/sec:   12.3 | eval time:    8.2 sec | output: 
    {'AP': np.float32(0.2243549),
     'AP50': np.float32(0.47362357),
     'AP75': np.float32(0.17409162),
     'APl': np.float32(0.31782418),
     'APm': np.float32(0.08430918),
     'APs': np.float32(0.0),
     'ARl': np.float32(0.33479053),
     'ARm': np.float32(0.1496675),
     'ARmax1': np.float32(0.16859457),
     'ARmax10': np.float32(0.24553351),
     'ARmax100': np.float32(0.30370182),
     'ARs': np.float32(0.0),
     'box_loss': np.float32(0.005913249),
     'cls_loss': np.float32(0.3845843),
     'model_loss': np.float32(0.68024683),
     'steps_per_second': 12.26576721886985,
     'total_loss': np.float32(1.6434913),
     'validation_loss': np.float32(1.6434913)}
train | step:    700 | training until step 800...
train | step:    800 | steps/sec:    0.7 | output: 
    {'box_loss': np.float32(0.005563259),
     'cls_loss': np.float32(0.3610574),
     'learning_rate': np.float32(0.009549147),
     'model_loss': np.float32(0.63922036),
     'total_loss': np.float32(1.6007898),
     'training_loss': np.float32(1.6007898)}
saved checkpoint to ./trained_model/ckpt-800.
 eval | step:    800 | running 100 steps of evaluation...
2026-11-05 12:54:27.111339: W tensorflow/core/framework/dataset.cc:993] Input of GeneratorDatasetOp::Dataset will not be optimized because the dataset does not implement the AsGraphDefInternal() method needed to apply optimizations.
creating index...
index created!
creating index...
index created!
Running per image evaluation...
Evaluate annotation type *bbox*
DONE (t=1.22s).
Accumulating evaluation results...
DONE (t=0.05s).
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.247
 Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets=100 ] = 0.532
 Average Precision  (AP) @[ IoU=0.75      | area=   all | maxDets=100 ] = 0.175
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.000
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.172
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.310
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=  1 ] = 0.184
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets= 10 ] = 0.276
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.336
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.000
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.240
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.335
 eval | step:    800 | steps/sec:   12.5 | eval time:    8.0 sec | output: 
    {'AP': np.float32(0.24669671),
     'AP50': np.float32(0.532271),
     'AP75': np.float32(0.1751407),
     'APl': np.float32(0.3099662),
     'APm': np.float32(0.17167515),
     'APs': np.float32(0.0),
     'ARl': np.float32(0.33469945),
     'ARm': np.float32(0.24037406),
     'ARmax1': np.float32(0.1835396),
     'ARmax10': np.float32(0.27565598),
     'ARmax100': np.float32(0.33592826),
     'ARs': np.float32(0.0),
     'box_loss': np.float32(0.0053206226),
     'cls_loss': np.float32(0.3642629),
     'model_loss': np.float32(0.630294),
     'steps_per_second': 12.492002079466785,
     'total_loss': np.float32(1.5905195),
     'validation_loss': np.float32(1.5905195)}
train | step:    800 | training until step 900...
train | step:    900 | steps/sec:    0.8 | output: 
    {'box_loss': np.float32(0.0054678083),
     'cls_loss': np.float32(0.35867885),
     'learning_rate': np.float32(0.002447176),
     'model_loss': np.float32(0.63206923),
     'total_loss': np.float32(1.5915761),
     'training_loss': np.float32(1.5915761)}
saved checkpoint to ./trained_model/ckpt-900.
 eval | step:    900 | running 100 steps of evaluation...
2026-11-05 12:56:39.591625: W tensorflow/core/framework/dataset.cc:993] Input of GeneratorDatasetOp::Dataset will not be optimized because the dataset does not implement the AsGraphDefInternal() method needed to apply optimizations.
creating index...
index created!
creating index...
index created!
Running per image evaluation...
Evaluate annotation type *bbox*
DONE (t=1.20s).
Accumulating evaluation results...
DONE (t=0.05s).
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.261
 Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets=100 ] = 0.535
 Average Precision  (AP) @[ IoU=0.75      | area=   all | maxDets=100 ] = 0.199
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.000
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.177
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.305
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=  1 ] = 0.189
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets= 10 ] = 0.286
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.350
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.000
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.245
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.328
 eval | step:    900 | steps/sec:   12.0 | eval time:    8.3 sec | output: 
    {'AP': np.float32(0.26075608),
     'AP50': np.float32(0.53487736),
     'AP75': np.float32(0.19903305),
     'APl': np.float32(0.3049371),
     'APm': np.float32(0.17686819),
     'APs': np.float32(0.0),
     'ARl': np.float32(0.32795992),
     'ARm': np.float32(0.24476309),
     'ARmax1': np.float32(0.1894308),
     'ARmax10': np.float32(0.28552735),
     'ARmax100': np.float32(0.34963626),
     'ARs': np.float32(0.0),
     'box_loss': np.float32(0.004968071),
     'cls_loss': np.float32(0.35815418),
     'model_loss': np.float32(0.6065577),
     'steps_per_second': 12.024829674211228,
     'total_loss': np.float32(1.5655729),
     'validation_loss': np.float32(1.5655729)}
train | step:    900 | training until step 1000...
train | step:   1000 | steps/sec:    0.8 | output: 
    {'box_loss': np.float32(0.0053327666),
     'cls_loss': np.float32(0.3543621),
     'learning_rate': np.float32(0.0),
     'model_loss': np.float32(0.6210004),
     'total_loss': np.float32(1.5798652),
     'training_loss': np.float32(1.5798652)}
saved checkpoint to ./trained_model/ckpt-1000.
 eval | step:   1000 | running 100 steps of evaluation...
2026-11-05 12:58:52.781736: W tensorflow/core/framework/dataset.cc:993] Input of GeneratorDatasetOp::Dataset will not be optimized because the dataset does not implement the AsGraphDefInternal() method needed to apply optimizations.
creating index...
index created!
creating index...
index created!
Running per image evaluation...
Evaluate annotation type *bbox*
DONE (t=1.19s).
Accumulating evaluation results...
DONE (t=0.05s).
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.262
 Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets=100 ] = 0.544
 Average Precision  (AP) @[ IoU=0.75      | area=   all | maxDets=100 ] = 0.206
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.000
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.177
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.322
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=  1 ] = 0.186
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets= 10 ] = 0.286
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.349
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.000
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.246
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.342
 eval | step:   1000 | steps/sec:   12.5 | eval time:    8.0 sec | output: 
    {'AP': np.float32(0.26242572),
     'AP50': np.float32(0.5444919),
     'AP75': np.float32(0.20562716),
     'APl': np.float32(0.32166955),
     'APm': np.float32(0.17748503),
     'APs': np.float32(0.0),
     'ARl': np.float32(0.34189436),
     'ARm': np.float32(0.24592684),
     'ARmax1': np.float32(0.18626818),
     'ARmax10': np.float32(0.2857377),
     'ARmax100': np.float32(0.34856772),
     'ARs': np.float32(0.0),
     'box_loss': np.float32(0.004878935),
     'cls_loss': np.float32(0.3527283),
     'model_loss': np.float32(0.5966751),
     'steps_per_second': 12.542240649135017,
     'total_loss': np.float32(1.5554805),
     'validation_loss': np.float32(1.5554805)}
 eval | step:   1000 | running 100 steps of evaluation...
creating index...
index created!
creating index...
index created!
Running per image evaluation...
Evaluate annotation type *bbox*
DONE (t=1.20s).
Accumulating evaluation results...
DONE (t=0.05s).
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.262
 Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets=100 ] = 0.544
 Average Precision  (AP) @[ IoU=0.75      | area=   all | maxDets=100 ] = 0.206
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.000
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.177
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.322
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=  1 ] = 0.186
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets= 10 ] = 0.286
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.349
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.000
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.246
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.342
 eval | step:   1000 | steps/sec:   12.7 | eval time:    7.9 sec | output: 
    {'AP': np.float32(0.26242572),
     'AP50': np.float32(0.5444919),
     'AP75': np.float32(0.20562716),
     'APl': np.float32(0.32166955),
     'APm': np.float32(0.17748503),
     'APs': np.float32(0.0),
     'ARl': np.float32(0.34189436),
     'ARm': np.float32(0.24592684),
     'ARmax1': np.float32(0.18626818),
     'ARmax10': np.float32(0.2857377),
     'ARmax100': np.float32(0.34856772),
     'ARs': np.float32(0.0),
     'box_loss': np.float32(0.004878935),
     'cls_loss': np.float32(0.3527283),
     'model_loss': np.float32(0.5966751),
     'steps_per_second': 12.706852350594273,
     'total_loss': np.float32(1.5554805),
     'validation_loss': np.float32(1.5554805)}

Load logs in tensorboard.

%load_ext tensorboard
%tensorboard --logdir './trained_model/'

Saving and exporting the trained model.

The keras.Model object returned by train_lib.run_experiment expects the data to be normalized by the dataset loader using the same mean and variance statiscics in preprocess_ops.normalize_image(image, offset=MEAN_RGB, scale=STDDEV_RGB). This export function handles those details, so you can pass tf.uint8 images and get the correct results.

export_saved_model_lib.export_inference_graph(
    input_type='image_tensor',
    batch_size=1,
    input_image_size=[HEIGHT, WIDTH],
    params=exp_config,
    checkpoint_path=tf.train.latest_checkpoint(model_dir),
    export_dir=export_dir)
WARNING:tensorflow:Skipping full serialization of TF-Keras layer <official.vision.modeling.retinanet_model.RetinaNetModel object at 0x7fcc08297b90>, because it is not built.
WARNING:tensorflow:Skipping full serialization of TF-Keras layer <official.vision.modeling.retinanet_model.RetinaNetModel object at 0x7fcc08297b90>, because it is not built.
WARNING:tensorflow:Skipping full serialization of TF-Keras layer <official.vision.modeling.layers.detection_generator.MultilevelDetectionGenerator object at 0x7fcc08170750>, because it is not built.
WARNING:tensorflow:Skipping full serialization of TF-Keras layer <official.vision.modeling.layers.detection_generator.MultilevelDetectionGenerator object at 0x7fcc08170750>, because it is not built.
INFO:tensorflow:Assets written to: ./exported_model/assets
INFO:tensorflow:Assets written to: ./exported_model/assets

Inference from trained model

def load_image_into_numpy_array(path):
  """Load an image from file into a numpy array.

  Puts image into numpy array to feed into tensorflow graph.
  Note that by convention we put it into a numpy array with shape
  (height, width, channels), where channels=3 for RGB.

  Args:
    path: the file path to the image

  Returns:
    uint8 numpy array with shape (img_height, img_width, 3)
  """
  image = None
  if(path.startswith('http')):
    response = urlopen(path)
    image_data = response.read()
    image_data = BytesIO(image_data)
    image = Image.open(image_data)
  else:
    image_data = tf.io.gfile.GFile(path, 'rb').read()
    image = Image.open(BytesIO(image_data))

  (im_width, im_height) = image.size
  return np.array(image.getdata()).reshape(
      (1, im_height, im_width, 3)).astype(np.uint8)



def build_inputs_for_object_detection(image, input_image_size):
  """Builds Object Detection model inputs for serving."""
  image, _ = resize_and_crop_image(
      image,
      input_image_size,
      padded_size=input_image_size,
      aug_scale_min=1.0,
      aug_scale_max=1.0)
  return image

Visualize test data.

num_of_examples = 3

test_ds = tf.data.TFRecordDataset(
    './bccd_coco_tfrecords/test-00000-of-00001.tfrecord').take(
        num_of_examples)
show_batch(test_ds, num_of_examples)

png

Importing SavedModel.

imported = tf.saved_model.load(export_dir)
model_fn = imported.signatures['serving_default']

Visualize predictions.

input_image_size = (HEIGHT, WIDTH)
plt.figure(figsize=(20, 20))
min_score_thresh = 0.30 # Change minimum score for threshold to see all bounding boxes confidences.

for i, serialized_example in enumerate(test_ds):
  plt.subplot(1, 3, i+1)
  decoded_tensors = tf_ex_decoder.decode(serialized_example)
  image = build_inputs_for_object_detection(decoded_tensors['image'], input_image_size)
  image = tf.expand_dims(image, axis=0)
  image = tf.cast(image, dtype = tf.uint8)
  image_np = image[0].numpy()
  result = model_fn(image)
  visualization_utils.visualize_boxes_and_labels_on_image_array(
      image_np,
      result['detection_boxes'][0].numpy(),
      result['detection_classes'][0].numpy().astype(int),
      result['detection_scores'][0].numpy(),
      category_index=category_index,
      use_normalized_coordinates=False,
      max_boxes_to_draw=200,
      min_score_thresh=min_score_thresh,
      agnostic_mode=False,
      instance_masks=None,
      line_thickness=4)
  plt.imshow(image_np)
  plt.axis('off')

plt.show()

png