def initialize_dataset(self, tf_record, batch_size, shuffle_buffer=512):
"""
Initialize and prepare TFRecord dataset for training.
Args:
tf_record: TFRecord file.
batch_size: int, training batch size
shuffle_buffer: Buffer size for shuffling dataset.
Returns:
dataset.
"""
dataset = read_tfr(
tf_record, self.classes_file, get_feature_map(), self.max_boxes
)
dataset = dataset.shuffle(shuffle_buffer)
dataset = dataset.batch(batch_size)
dataset = dataset.map(
lambda x, y: (
transform_images(x, self.input_shape[0]),
transform_targets(
y, self.anchors, self.masks, self.input_shape[0]
),
)
)
dataset = dataset.prefetch(buffer_size=tf.data.experimental.AUTOTUNE)
return dataset
def predict_image(self, image_data, features):
"""
Make predictions on a single image from the TFRecord.
Args:
image_data: image as numpy array
features: features of the TFRecord.
Returns:
pandas DataFrame with detection data.
"""
image_path = bytes.decode(features['image_path'].numpy())
image_name = os.path.basename(image_path)
image = tf.expand_dims(image_data, 0)
resized = transform_images(image, self.input_shape[0])
outs = self.inference_model(resized)
adjusted = cv2.cvtColor(image_data.numpy(), cv2.COLOR_RGB2BGR)
result = (
get_detection_data(adjusted, image_name, outs, self.class_names),
image_name,
)
return result
def detect_image(self, image_data, image_name):
"""
Given an image, get detections for the image.
Args:
image_data: image as numpy array/tf.Tensor
image_name: str, name of the image
Returns:
pandas DataFrame with detections found.
"""
image = tf.expand_dims(image_data, 0)
resized = transform_images(image, self.input_shape[0])
out = self.inference_model.predict(resized)
if isinstance(image_data, np.ndarray):
adjusted = cv2.cvtColor(image_data, cv2.COLOR_RGB2BGR)
else:
adjusted = cv2.cvtColor(image_data.numpy(), cv2.COLOR_RGB2BGR)
detections = get_detection_data(
adjusted,
image_name,
out,
self.class_names,
)
return detections, adjusted