def build_model(self): """Build an instance of Yolo.""" model_base_cfg = self.task_config.model l2_weight_decay = self.task_config.weight_decay / 2.0 input_size = model_base_cfg.input_size.copy() input_specs = tf.keras.layers.InputSpec(shape=[None] + input_size) l2_regularizer = (tf.keras.regularizers.l2(l2_weight_decay) if l2_weight_decay else None) model, losses = factory.build_yolo(input_specs, model_base_cfg, l2_regularizer) # save for later usage within the task. self._loss_fn = losses self._model = model return model
def create_yolo_export_module( params: cfg.ExperimentConfig, input_type: str, batch_size: int, input_image_size: List[int], num_channels: int = 3) -> export_base.ExportModule: """Creates YOLO export module.""" input_signature = export_utils.get_image_input_signatures( input_type, batch_size, input_image_size, num_channels) input_specs = tf.keras.layers.InputSpec(shape=[batch_size] + input_image_size + [num_channels]) model, _ = yolo_factory.build_yolo(input_specs=input_specs, model_config=params.task.model, l2_regularization=None) def preprocess_fn(inputs): image_tensor = export_utils.parse_image(inputs, input_type, input_image_size, num_channels) # If input_type is `tflite`, do not apply image preprocessing. if input_type == 'tflite': return image_tensor def preprocess_image_fn(inputs): image = tf.cast(inputs, dtype=tf.float32) image = image / 255. (image, image_info) = yolo_model_fn.letterbox( image, input_image_size, letter_box=params.task.validation_data.parser.letter_box) return image, image_info images_spec = tf.TensorSpec(shape=input_image_size + [3], dtype=tf.float32) image_info_spec = tf.TensorSpec(shape=[4, 2], dtype=tf.float32) images, image_info = tf.nest.map_structure( tf.identity, tf.map_fn(preprocess_image_fn, elems=image_tensor, fn_output_signature=(images_spec, image_info_spec), parallel_iterations=32)) return images, image_info def inference_steps(inputs, model): images, image_info = inputs detection = model(images, training=False) detection['bbox'] = yolo_model_fn.undo_info( detection['bbox'], detection['num_detections'], image_info, expand=False) final_outputs = { 'detection_boxes': detection['bbox'], 'detection_scores': detection['confidence'], 'detection_classes': detection['classes'], 'num_detections': detection['num_detections'] } return final_outputs export_module = export_base.ExportModule(params, model=model, input_signature=input_signature, preprocessor=preprocess_fn, inference_step=inference_steps) return export_module