{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Introduction\n\n[Object detection API](https://tensorflow-object-detection-api-tutorial.readthedocs.io/en/latest/) is a tensorflow-based library for object detectio tasks. Following [the docs](https://tensorflow-object-detection-api-tutorial.readthedocs.io/en/latest/training.html) you should be able to train your own detector without problems. Furthermore, there is a well written [public notebook](https://www.kaggle.com/sreevishnudamodaran/vbd-efficientdet-tf2-object-detection-api) that uses OD API for this challenge.\n\nThe API itself provides a python script for both training and evaluation of a detection model, but this approach is not very suitable for IPython notebooks and is not easy to customize. For this reason I've choosen a slightly different approach to use this API by looking at their codebase and rewriting the training loop by myself. Most of the code is indeed copied from OD API with some minor modifications to make it work in kaggle notebooks.\n\nI'm sorry if the code isn't very clear, this is my first attempt with object detection API too."},{"metadata":{},"cell_type":"markdown","source":"# Install Object Detection API\n\n**NOTE:** I decided to use commit `3f6fe2aa410d901aae8829597a65d084bffc20d3` as it does not require tensorflow version `2.4.0` (that causes CUDA version mismatch due to some recent commit)."},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"%%capture\n!git clone https://github.com/tensorflow/models.git\n\n%cd models/research/\n!git reset --hard 3f6fe2aa410d901aae8829597a65d084bffc20d3\n\n!protoc object_detection/protos/*.proto --python_out=.\n\n!cp object_detection/packages/tf2/setup.py .\n!python -m pip install . \n%cd /kaggle/working","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Setup directories\n\n* Organise workspace/training files: in standard OD API approach you need to setup your working directories following a specific tree. For convenience, I did the same here:\n\n```\n./\n├─ annotations/\n    └─ label_map.pbtxt\n├─ models/\n    └─ pipeline.config\n└─ pre-trained-models/*/\n    ├─ checkpoint/\n    ├─ saved_model/\n    └─ pipeline.config\n```\n\n* Download pre-trained EfficientDet-D0 from [TF Model Zoo](https://github.com/tensorflow/models/blob/master/research/object_detection/g3doc/tf2_detection_zoo.md)"},{"metadata":{"trusted":true},"cell_type":"code","source":"BASE_DIR = 'chest-x-ray-detection'\nMODEL_PATH = 'efficientdet_d0_coco17_tpu-32'","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"%%capture\n!rm -r {BASE_DIR}\n!mkdir {BASE_DIR}\n!mkdir {BASE_DIR}/pre-trained-models/\n!mkdir {BASE_DIR}/annotations\n!mkdir {BASE_DIR}/models\n!mkdir {BASE_DIR}/models/efficientdet/\n\n!wget http://download.tensorflow.org/models/object_detection/tf2/20200711/{MODEL_PATH}.tar.gz\n\n!tar -xvzf {MODEL_PATH}.tar.gz\n!rm {MODEL_PATH}.tar.gz\n!mv {MODEL_PATH} {BASE_DIR}/pre-trained-models/\n!mv {BASE_DIR}/pre-trained-models/{MODEL_PATH}/pipeline.config {BASE_DIR}/models/efficientdet/pipeline.config","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Import libraries"},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"import pathlib, cv2, os, time, functools\nimport numpy as np\nimport pandas as pd\n\nimport matplotlib.pyplot as plt\n%matplotlib inline\n\nimport tensorflow as tf\n\nfrom google.protobuf import text_format","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Object Detection API internal libraries"},{"metadata":{"trusted":true},"cell_type":"code","source":"from object_detection import inputs\n\nfrom object_detection.model_lib_v2 import eager_train_step\nfrom object_detection.model_lib_v2 import eager_eval_loop\nfrom object_detection.model_lib_v2 import load_fine_tune_checkpoint\nfrom object_detection.model_lib_v2 import get_filepath\nfrom object_detection.model_lib_v2 import clean_temporary_directories\n\nfrom object_detection.protos import pipeline_pb2\n\nfrom object_detection.utils import label_map_util\nfrom object_detection.utils import visualization_utils as viz_utils\nfrom object_detection.utils import config_util\n\nfrom object_detection.builders import dataset_builder\nfrom object_detection.builders import image_resizer_builder\nfrom object_detection.builders import model_builder\nfrom object_detection.builders import preprocessor_builder\n\nfrom object_detection.core import standard_fields as fields\n\nfrom object_detection.exporter_lib_v2 import DetectionInferenceModule","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Setup notebook"},{"metadata":{"trusted":true},"cell_type":"code","source":"MODEL_DIR = BASE_DIR + '/models/efficientdet/'\nPIPELINE_PATH = MODEL_DIR + 'pipeline.config'\nLABEL_MAP_PATH = BASE_DIR + '/annotations/label_map.pbtxt'\nOUTPUT_MODEL_DIR = '/kaggle/working/saved_model'","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"input_path = pathlib.Path('/kaggle/input/chest-xray-detection-512x512-groupkfold-tfrec')\n\n!cp {input_path}/label_map.pbtxt {LABEL_MAP_PATH}\n\nDS_PATH = str(input_path)\nos.makedirs(OUTPUT_MODEL_DIR, exist_ok=True)","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"plt.rcParams['axes.grid'] = False\nplt.rcParams['xtick.labelsize'] = False\nplt.rcParams['ytick.labelsize'] = False\nplt.rcParams['xtick.top'] = False\nplt.rcParams['xtick.bottom'] = False\nplt.rcParams['ytick.left'] = False\nplt.rcParams['ytick.right'] = False\nplt.rcParams['figure.figsize'] = [12, 12]","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"def seed_everything(seed=0):\n    np.random.seed(seed)\n    tf.random.set_seed(seed)\n    os.environ['PYTHONHASHSEED'] = str(seed)\n    os.environ['TF_DETERMINISTIC_OPS'] = '1'\n\nseed = 2020\nseed_everything(seed)","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"gpus = tf.config.experimental.list_physical_devices('GPU')\nif gpus:\n    try:\n        tf.config.experimental.set_visible_devices(gpus[0], 'GPU')\n        logical_gpus = tf.config.experimental.list_logical_devices('GPU')\n        print(len(gpus), \"Physical GPUs,\", len(logical_gpus), \"Logical GPU\")\n        strategy = tf.distribute.MirroredStrategy(devices=[\"GPU:0\"])\n    except RuntimeError as e:\n        gpu = None","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-input":false,"trusted":true},"cell_type":"code","source":"NUM_CLASSES = 14\nPER_REPLICA_BATCH_SIZE = 2\ntry:\n    REPLICAS = strategy.num_replicas_in_sync\nexcept:\n    REPLICAS = 1\n    \nBATCH_SIZE = PER_REPLICA_BATCH_SIZE * REPLICAS\n\nfold = 0\nN_FOLDS = 5\n\nSCORE_THRESHOLD = 0.5","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Read data\n\nI created a custom version of the dataset stored in tfrecord files (private at the moment). I'm not planning to make the dataset public, but I've done basic pre-processing steps:\n\n* voi-lut\n* monochrome correction\n* resizing to `512x512`\n* histogram equalization\n\nI've also splitted data by patient id in `5 folds` using (Multi-Class Stratified) GroupKFold."},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df = pd.read_csv(input_path / 'train.csv')\ntrain_df.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"category_index = label_map_util.create_category_index_from_labelmap(\n    LABEL_MAP_PATH,\n    use_display_name=True\n)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"TRAIN_DATASET = tf.io.gfile.glob(DS_PATH + f'/fold_[^{fold + 1}].tfrecord')\nTEST_DATASET = tf.io.gfile.glob(DS_PATH + f'/fold_{fold + 1}.tfrecord')    \n\nct_train = len(train_df['image_id'][train_df['fold'] != fold + 1].unique())  / BATCH_SIZE\nct_test = len(train_df['image_id'][train_df['fold'] == fold + 1].unique()) / BATCH_SIZE","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Show samples"},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"def plot_img_with_boxes(image, classes, boxes, scores=None, axis=None, plot=True):\n    if scores is None:\n        scores = np.ones(len(classes))\n        \n    image_with_detections = image.copy()\n    \n    viz_utils.visualize_boxes_and_labels_on_image_array(\n          image_with_detections,\n          boxes,\n          classes,\n          scores,\n          category_index,\n          use_normalized_coordinates=True,\n          max_boxes_to_draw=100,\n          min_score_thresh=SCORE_THRESHOLD,\n          agnostic_mode=False)\n    \n    if plot:\n        if axis is None:\n            plt.figure(figsize=(12,12))\n            plt.imshow(image_with_detections)\n            plt.show()\n        else:\n            axis.imshow(image_with_detections)\n    else:\n        return image_with_detections","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"feature_description = {\n    'image/height': tf.io.FixedLenFeature([], tf.int64),\n    'image/width': tf.io.FixedLenFeature([], tf.int64),\n    'image/filename': tf.io.FixedLenFeature([], tf.string),\n    'image/source_id': tf.io.FixedLenFeature([], tf.string),\n    'image/encoded': tf.io.FixedLenFeature([], tf.string),\n    'image/format': tf.io.FixedLenFeature([], tf.string),\n    'image/object/bbox/xmin': tf.io.FixedLenSequenceFeature([], tf.float32, True),\n    'image/object/bbox/xmax': tf.io.FixedLenSequenceFeature([], tf.float32, True),\n    'image/object/bbox/ymin': tf.io.FixedLenSequenceFeature([], tf.float32, True),\n    'image/object/bbox/ymax': tf.io.FixedLenSequenceFeature([], tf.float32, True),\n    'image/object/class/text': tf.io.FixedLenSequenceFeature([], tf.string, True),\n    'image/object/class/label': tf.io.FixedLenSequenceFeature([], tf.int64, True)\n}\n\ndef parse_image_sample(example_proto):\n    return tf.io.parse_single_example(example_proto,\n                                      feature_description)\n\nraw_image_dataset = tf.data.TFRecordDataset(TEST_DATASET)\nparsed_image_dataset = raw_image_dataset.map(parse_image_sample)\niterator = iter(parsed_image_dataset)","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"%matplotlib inline\n\nN = 16\nfig, ax = plt.subplots(int(np.sqrt(N)), int(np.sqrt(N)), figsize=(12,12))\nax = ax.flatten()\n\nfor idx in range(N):\n    image_features = next(iterator)\n    image_raw = image_features['image/encoded']\n    image = tf.image.decode_jpeg(image_raw).numpy()\n    classes = image_features['image/object/class/label'].numpy()\n    boxes = np.stack([\n        image_features['image/object/bbox/xmin'],\n        image_features['image/object/bbox/ymin'],\n        image_features['image/object/bbox/xmax'],\n        image_features['image/object/bbox/ymax'],\n    ], -1)\n\n    plot_img_with_boxes(image, \n                        classes, \n                        boxes,\n                        axis=ax[idx])\n    \nfig.show()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Model configuration\n\nHere I override default configurations of the pre-trained model that are specific for COCO dataset. This step is equivalent to writing the `pipeline.config` file by hand or by looking for strings through regular expressions, but I find this way more clean and straightforward.\n\nNotice that differently from the standard way of using OD API, in this case I'm not using the `pipeline.config` file, but I'm loading configurations into a dictionary that I'm going to use later."},{"metadata":{"trusted":true},"cell_type":"code","source":"configs = config_util.get_configs_from_pipeline_file(PIPELINE_PATH)\n\nconfigs['model'].ssd.num_classes = NUM_CLASSES\n\nconfigs['train_config'].sync_replicas = True if REPLICAS > 1 else False\nconfigs['train_config'].replicas_to_aggregate = REPLICAS\nconfigs['train_config'].batch_size = BATCH_SIZE\nconfigs['train_config'].data_augmentation_options.pop(1)\n\nconfigs['train_config'].fine_tune_checkpoint = (\n    BASE_DIR + f'/pre-trained-models/{MODEL_PATH}/checkpoint/ckpt-0'\n)\nconfigs['train_config'].fine_tune_checkpoint_type = \"detection\"\n\nconfigs['train_input_config'].label_map_path = LABEL_MAP_PATH\nconfigs['train_input_config'].tf_record_input_reader.input_path[:] = TRAIN_DATASET\nconfigs['train_input_config'].load_multiclass_scores = True\n\nconfigs['eval_config'].batch_size = 1\nconfigs['eval_config'].metrics_set[:] = ''\nconfigs['eval_config'].metrics_set.append('pascal_voc_detection_metrics')\n\nconfigs['eval_input_config'].label_map_path = LABEL_MAP_PATH\nconfigs['eval_input_config'].tf_record_input_reader.input_path[:] = TEST_DATASET\nconfigs['eval_input_config'].load_multiclass_scores = True\n\nconfig_util.save_pipeline_config(config_util.create_pipeline_proto_from_configs(configs),\n                                 PIPELINE_PATH.replace('pipeline.config', ''))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model_config = configs['model']\ntrain_config = configs['train_config']\ntrain_input_config = configs['train_input_config']\neval_config = configs['eval_config']\neval_input_config = configs['eval_input_configs'][0]","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Build model\n\nHere I build my detector using `model_builder` utils. Notice that `build_model` fuction also overrides the preprocessing function used by the feature extractor backbone (`feature_extractor.preprocess`). This function is used both during training and inference to pre-process the data before feeding them to the detector. The [default pre-processing](https://github.com/tensorflow/models/blob/master/research/object_detection/models/ssd_efficientnet_bifpn_feature_extractor.py#L185)  provided by OD API is channel-wise normalization by ImageNet mean/std."},{"metadata":{"trusted":true},"cell_type":"code","source":"def preprocess_fn(inputs):\n    return inputs / 255.0\n\ndef build_model():\n    detection_model = model_builder._build_ssd_model(ssd_config=model_config.ssd,\n                                                     is_training=True,\n                                                     add_summaries=False)\n\n    detection_model._feature_extractor.preprocess = preprocess_fn\n    \n    return detection_model\n\ntry:\n    with strategy.scope():\n        detection_model = build_model()\nexcept:\n    detection_model = build_model()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Training utils"},{"metadata":{"trusted":true},"cell_type":"code","source":"SHAPE = (512,512)\nlearning_rate = 1e-4\n\nEPOCHS = 1\nSTEPS_PER_EPOCH = int(ct_train)\nNUM_TRAIN_STEPS = int(STEPS_PER_EPOCH * EPOCHS)\ntrain_steps = NUM_TRAIN_STEPS\n\nRUN_EVAL = True\nMONITOR_METRIC = 'PascalBoxes_Precision/mAP@0.5IOU'\nES_PATIENCE = 5\n\nbest_metric_value = 0.0\nnot_improved = 0\nsteps_per_sec_list = []\n\nunpad_groundtruth_tensors = train_config.unpad_groundtruth_tensors\nadd_regularization_loss = train_config.add_regularization_loss\n\nclip_gradients_value = None\nif train_config.gradient_clipping_by_norm > 0:\n    clip_gradients_value = train_config.gradient_clipping_by_norm\n\nconfig_util.update_fine_tune_checkpoint_type(train_config)\nfine_tune_checkpoint_type = train_config.fine_tune_checkpoint_type\nfine_tune_checkpoint_version = train_config.fine_tune_checkpoint_version","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Dataset functions\n\nI used `dataset.build` ([code](https://github.com/tensorflow/models/blob/master/research/object_detection/builders/dataset_builder.py#L166)) as it is done in OD API codebase. From the docs:\n```\nBuilds a tf.data.Dataset by applying the `transform_input_data_fn` on all\n  records. Applies a padded batch to the resulting dataset.\n  Args:\n    input_reader_config: A input_reader_pb2.InputReader object.\n    batch_size: Batch size. If batch size is None, no batching is performed.\n    transform_input_data_fn: Function to apply transformation to all records,\n      or None if no extra decoding is required.\n    input_context: optional, A tf.distribute.InputContext object used to\n      shard filenames and compute per-replica batch_size when this function\n      is being called per-replica.\n    reduce_to_frame_fn: Function that extracts frames from tf.SequenceExample\n      type input data.\n  Returns:\n    A tf.data.Dataset based on the input_reader_config.\n```"},{"metadata":{},"cell_type":"markdown","source":"### Training dataset\n\nHere you can change `train_config.data_augmentation_options` in `pipeline.config` file to change data augmentation applied during training. Possible options are listed [here](https://github.com/tensorflow/models/blob/master/research/object_detection/protos/preprocessor.proto#L8)."},{"metadata":{"_kg_hide-input":false,"trusted":true},"cell_type":"code","source":"def train_dataset_fn(input_context):\n    def transform_input_data_fn(tensor_dict):\n        data_augmentation_options = [\n            preprocessor_builder.build(step)\n            for step in train_config.data_augmentation_options\n        ]\n        data_augmentation_fn = functools.partial(\n            inputs.augment_input_data,\n            data_augmentation_options=data_augmentation_options\n        )\n\n        image_resizer_config = model_config.ssd.image_resizer\n        image_resizer_fn = image_resizer_builder.build(image_resizer_config)\n        transform_data_fn = functools.partial(\n            inputs.transform_input_data, \n            model_preprocess_fn=detection_model.preprocess,\n            image_resizer_fn=image_resizer_fn,\n            num_classes=NUM_CLASSES,\n            data_augmentation_fn=data_augmentation_fn,\n            merge_multiple_boxes=False,\n            use_multiclass_scores=False\n        )\n\n        tensor_dict = inputs.pad_input_data_to_static_shapes(\n            tensor_dict=transform_data_fn(tensor_dict),\n            max_num_boxes=train_input_config.max_number_of_boxes,\n            num_classes=NUM_CLASSES,\n            spatial_image_shape=SHAPE\n        )\n\n        return (inputs._get_features_dict(tensor_dict, False),\n                inputs._get_labels_dict(tensor_dict))\n\n    train_input = dataset_builder.build(\n        train_input_config,\n        transform_input_data_fn=transform_input_data_fn,\n        batch_size=train_config.batch_size,\n        input_context=input_context,\n    )\n    train_input = train_input.repeat()    \n\n    return train_input","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Validation dataset"},{"metadata":{"_kg_hide-input":false,"trusted":true},"cell_type":"code","source":"def eval_dataset_fn(input_context):\n    def transform_input_data_fn(tensor_dict):\n        image_resizer_config = model_config.ssd.image_resizer\n        image_resizer_fn = image_resizer_builder.build(image_resizer_config)\n\n        transform_data_fn = functools.partial(\n            inputs.transform_input_data, \n            model_preprocess_fn=detection_model.preprocess,\n            image_resizer_fn=image_resizer_fn,\n            num_classes=NUM_CLASSES,\n            merge_multiple_boxes=False,\n            use_multiclass_scores=False,\n            retain_original_image=eval_config.retain_original_images,\n            retain_original_image_additional_channels=eval_config.retain_original_image_additional_channels\n        )\n\n        tensor_dict = inputs.pad_input_data_to_static_shapes(\n            tensor_dict=transform_data_fn(tensor_dict),\n            max_num_boxes=eval_input_config.max_number_of_boxes,\n            num_classes=NUM_CLASSES,\n            spatial_image_shape=SHAPE\n        )\n\n        return (inputs._get_features_dict(tensor_dict, False),\n                inputs._get_labels_dict(tensor_dict))\n\n    eval_input = dataset_builder.build(\n        eval_input_config,\n        transform_input_data_fn=transform_input_data_fn,\n        batch_size=eval_config.batch_size,\n        input_context=input_context,\n    )\n\n    return eval_input","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-input":false,"trusted":true},"cell_type":"code","source":"train_input = strategy.experimental_distribute_datasets_from_function(\n    train_dataset_fn\n)\n\ntrain_input_iter = iter(train_input)\n\neval_input = strategy.experimental_distribute_datasets_from_function(\n    eval_dataset_fn\n)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Show training samples"},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"%matplotlib inline\n\nN = 16\nfig, ax = plt.subplots(int(np.sqrt(N)), int(np.sqrt(N)), figsize=(12,12))\nax = ax.flatten()\n\ntrain_dataset = train_dataset_fn(None).unbatch().batch(1)\ntrain_iter = iter(train_dataset)\n\nfor idx in range(N):\n    features, labels = next(train_iter)\n    image = features['image'][0].numpy()\n    n_boxes = labels['num_groundtruth_boxes'][0].numpy()\n    boxes = labels['groundtruth_boxes'][0, :n_boxes, :].numpy()\n    classes = labels['groundtruth_classes'][0, :n_boxes, :].numpy()\n    classes = np.argmax(classes, axis=-1) + 1\n    \n    plot_img_with_boxes((image*255).astype('uint8'), \n                        classes, \n                        boxes,\n                        axis=ax[idx])\n    \nfig.show()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Training"},{"metadata":{},"cell_type":"markdown","source":"## Load weights\n\nFirst I load the model checkpoint from `pipeline.config` file (i.e. the saved model pre-trained on COCO dataset)."},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"with strategy.scope():\n    global_step = tf.Variable(0,\n                              trainable=False,\n                              dtype=tf.compat.v2.dtypes.int64,\n                              name='global_step',\n                              aggregation=tf.compat.v2.VariableAggregation.ONLY_FIRST_REPLICA)\n    \n    checkpointed_step = int(global_step.value())\n    logged_step = int(global_step.value())\n    total_loss = 0\n\n    if train_config.fine_tune_checkpoint:\n        load_fine_tune_checkpoint(detection_model,\n                                  train_config.fine_tune_checkpoint,\n                                  fine_tune_checkpoint_type,\n                                  fine_tune_checkpoint_version,\n                                  train_input,\n                                  unpad_groundtruth_tensors)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Then I load the last saved checkpoint from `MODEL_DIR` (if any), in order to resume the training process if any checkpoint has been saved in `MODEL_DIR`."},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"with strategy.scope():\n    if callable(learning_rate):\n        learning_rate_fn = learning_rate\n    else:\n        learning_rate_fn = lambda: learning_rate\n\n    optimizer = tf.keras.optimizers.Adam(learning_rate=learning_rate)\n    \n    ckpt = tf.compat.v2.train.Checkpoint(step=global_step,\n                                         model=detection_model,\n                                         optimizer=optimizer)\n\n    manager_dir = get_filepath(strategy, MODEL_DIR)\n\n    manager = tf.compat.v2.train.CheckpointManager(ckpt,\n                                                   manager_dir,\n                                                   max_to_keep=1)\n\n    latest_checkpoint = tf.train.latest_checkpoint(MODEL_DIR)\n    ckpt.restore(latest_checkpoint).expect_partial()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Custom training loop\n\nTo implement the training loop i used `eager_train_step`  from `model_lib_v2` ([code](https://github.com/tensorflow/models/blob/31e86e86c1e7f4154819e1c52ea0c51b287c2c70/research/object_detection/model_lib_v2.py#L146)), but you can rewrite if on your own as it is essentially a standard [TF-2 custom training loop](https://www.tensorflow.org/guide/keras/writing_a_training_loop_from_scratch), except for the fact that I use COCO evaluator to compute mean average precision after each epoch."},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"with strategy.scope():\n    def train_step_fn(features, labels):\n        loss = eager_train_step(detection_model,\n                                features,\n                                labels,\n                                unpad_groundtruth_tensors,\n                                optimizer,\n                                learning_rate=learning_rate_fn(),\n                                add_regularization_loss=add_regularization_loss,\n                                clip_gradients_value=clip_gradients_value,\n                                global_step=global_step,\n                                num_replicas=REPLICAS)\n        global_step.assign_add(1)\n        return loss\n\n    def _sample_and_train(strategy, train_step_fn, data_iterator):\n        features, labels = data_iterator.next()\n        per_replica_losses = strategy.run(train_step_fn, \n                                          args=(features, labels))\n        return strategy.reduce(tf.distribute.ReduceOp.SUM,\n                               per_replica_losses, axis=None)\n\n    @tf.function\n    def _dist_train_step(data_iterator):\n        return _sample_and_train(strategy, \n                                 train_step_fn, \n                                 data_iterator)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Notice that in the following cell the training steps are performed within `strategy.scope()`, while the evaluation loop is performed outside. This is intended as differently from the training step, for evaluation i used directly `eager_eval_loop`. If you look at the [code](https://github.com/tensorflow/models/blob/31e86e86c1e7f4154819e1c52ea0c51b287c2c70/research/object_detection/model_lib_v2.py#L775) you will notice that this function use a similar approach to the one used in previous cell. Anyway the GPU usage when running evaluation is very low with respect to the training loop. For this reason, a more efficient and time-saving approach would be to write functions similar to those used for training also for running evaluation, at the cost of not using the mAP score but for instance the value of the loss (or other tensorflow metrics)."},{"metadata":{"trusted":true},"cell_type":"code","source":"last_step_time = time.time()\n\nfor _ in range(global_step.value(), train_steps):\n    with strategy.scope():\n        loss = _dist_train_step(train_input_iter)\n        time_taken = time.time() - last_step_time\n        last_step_time = time.time()\n        steps_per_sec = 1.0 / time_taken\n        steps_per_sec_list.append(steps_per_sec)\n        total_loss += loss\n        \n    if int(global_step.value()) % STEPS_PER_EPOCH == 0:\n        if not RUN_EVAL:\n            print('Epoch {} [ETA {:.2f}s] loss={:.3f}'.format(\n                  int(global_step.value()) // STEPS_PER_EPOCH,\n                  time_taken * STEPS_PER_EPOCH,\n                  total_loss / STEPS_PER_EPOCH))\n        else:\n            eval_global_step = tf.compat.v2.Variable(0, \n                                                     trainable=False,\n                                                     dtype=tf.compat.v2.dtypes.int64)\n\n            eval_metrics = eager_eval_loop(detection_model,\n                                           configs,\n                                           eval_input,\n                                           global_step=eval_global_step)\n\n            print('Epoch {} [ETA {:.2f}s] loss={:.3f} mAP@.5={:.3f}'.format(\n                  int(global_step.value()) // STEPS_PER_EPOCH,\n                  time_taken * STEPS_PER_EPOCH,\n                  total_loss / STEPS_PER_EPOCH,\n                  eval_metrics[MONITOR_METRIC]))\n\n            if eval_metrics[MONITOR_METRIC] > best_metric_value:\n                best_metric_value = eval_metrics[MONITOR_METRIC]\n                manager.save()\n                not_improved = 0\n            else:\n                not_improved += 1\n\n            if not_improved >= ES_PATIENCE:\n                print(f\"Early stopping at epoch {int(global_step.value()) // STEPS_PER_EPOCH}\")\n                break\n            \n        total_loss = 0\n\nclean_temporary_directories(strategy, manager_dir)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Export model for inference\n\nWe first load the latest checkpoint saved during training, then we export the whole detection model (pre-processing and post-processing included) in TF2 OD-API style."},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"class DetectionFromImageModule(DetectionInferenceModule):\n    def __init__(self, detection_model):\n        \n        sig = [tf.TensorSpec(shape=[1, None, None, 3],\n                             dtype=tf.uint8,\n                             name='input_tensor')]\n\n        def call_func(input_tensor):\n            return self._run_inference_on_images(input_tensor)\n\n        self.__call__ = tf.function(call_func, input_signature=sig)\n\n        super(DetectionFromImageModule, self).__init__(detection_model)\n        \n    def _run_inference_on_images(self, image, **kwargs):\n        label_id_offset = 1\n        image = tf.cast(image, tf.float32)\n        image, shapes = self._model.preprocess(image)\n        prediction_dict = self._model.predict(image, shapes, **kwargs)\n        detections = self._model.postprocess(prediction_dict, shapes)\n        classes_field = fields.DetectionResultFields.detection_classes\n        classes = tf.cast(detections[classes_field], tf.float32)\n        detections[classes_field] = (classes + label_id_offset)\n\n        for key, val in detections.items():\n            detections[key] = tf.cast(val, tf.float32)\n\n        return detections","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"ckpt = tf.train.Checkpoint(model=detection_model)\nmanager = tf.train.CheckpointManager(ckpt, \n                                     MODEL_DIR,\n                                     max_to_keep=1)\n\nstatus = ckpt.restore(manager.latest_checkpoint).expect_partial()\n\ndetection_module = DetectionFromImageModule(detection_model)\nconcrete_function = detection_module.__call__.get_concrete_function()\nstatus.assert_existing_objects_matched()\n\nexported_checkpoint_manager = tf.train.CheckpointManager(ckpt, \n                                                         OUTPUT_MODEL_DIR, \n                                                         max_to_keep=1)\n\nexported_checkpoint_manager.save(checkpoint_number=0)\ntf.saved_model.save(detection_module,\n                    OUTPUT_MODEL_DIR + '/saved_model',\n                    signatures=concrete_function)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Inference from saved model\n\nHere we test the saved model by running it on few examples and compare detected boxes (coloured) with respect to ground-truth boxes (in black)."},{"metadata":{"trusted":true},"cell_type":"code","source":"detector = tf.saved_model.load('/kaggle/working/saved_model/saved_model/')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Plot detected boxes"},{"metadata":{"trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"%matplotlib inline\n\nN = 16\nfig, ax = plt.subplots(int(np.sqrt(N)), int(np.sqrt(N)), figsize=(12,12))\nax = ax.flatten()\n\nraw_image_dataset = tf.data.TFRecordDataset(TEST_DATASET)\nparsed_image_dataset = raw_image_dataset.map(parse_image_sample)\niterator = iter(parsed_image_dataset)\n\nfor idx in range(N):\n    image_features = next(iterator)\n    image_raw = image_features['image/encoded']\n    image = tf.image.decode_jpeg(image_raw)\n    gt_boxes = np.stack([\n        image_features['image/object/bbox/xmin'],\n        image_features['image/object/bbox/ymin'],\n        image_features['image/object/bbox/xmax'],\n        image_features['image/object/bbox/ymax'],\n    ], -1)\n    \n    out = detector(tf.expand_dims(image, 0))\n    \n    classes = out['detection_classes'].numpy()[0].astype('int')\n    scores = out['detection_scores'].numpy()[0]\n    boxes = out['detection_boxes'].numpy()[0]\n    boxes = np.stack([\n        boxes[:,1],\n        boxes[:,0],\n        boxes[:,3],\n        boxes[:,2]        \n    ], -1)\n    \n    image_with_gt = image.numpy()\n    \n    viz_utils.draw_bounding_boxes_on_image_array(image_with_gt, \n                                                 gt_boxes,\n                                                 color='black')\n    \n    plot_img_with_boxes(image_with_gt, \n                        classes, \n                        boxes,\n                        scores,\n                        axis=ax[idx])\n    \nfig.show()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Clean Environment"},{"metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"trusted":true},"cell_type":"code","source":"!rm -r /kaggle/working/models\n!rm -r /kaggle/working/chest-x-ray-detection","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}