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