src.models.keypoint_detector
1import keras 2import keras.ops as ops 3import numpy as np 4import matplotlib.pyplot as plt 5from tqdm import tqdm 6 7from src.utils import * 8 9 10class KeypointDetector(keras.layers.Layer): 11 def __init__(self, num_keypoints=10, keypoint_branch=None, jacobian_branch=None, **kwargs): 12 """ 13 Layer for detecting keypoints and local Jacobian transformations. 14 15 Args: 16 num_keypoints (int): Number of keypoints to detect. 17 keypoint_branch (keras.Model or None): Optional custom model for keypoint detection. 18 If None, a default convolutional branch will be used. 19 jacobian_branch (keras.Model or None): Optional custom model for Jacobian prediction. 20 If None, a default convolutional branch will be used. 21 **kwargs: Additional keyword arguments passed to the parent class. 22 """ 23 super().__init__(**kwargs) 24 self.num_keypoints = num_keypoints 25 26 if keypoint_branch is None: 27 self.keypoint_branch = keras.Sequential([ 28 keras.layers.Conv2D(32, 3, strides=2, activation='relu', padding='same'), 29 keras.layers.Conv2D(32, 3, strides=2, activation='relu', padding='same'), 30 keras.layers.Conv2D(num_keypoints, 3, strides=1, activation=None, padding='same'), 31 ]) 32 else: 33 self.keypoint_branch = keypoint_branch 34 35 if jacobian_branch is None: 36 self.jacobian_branch = keras.Sequential([ 37 keras.layers.Conv2D(32, 3, strides=2, activation='relu', padding='same'), 38 keras.layers.Conv2D(32, 3, strides=2, activation='relu', padding='same'), 39 keras.layers.Conv2D(4 * num_keypoints, 3, strides=1, activation=None, padding='same'), 40 ]) 41 else: 42 self.jacobian_branch = jacobian_branch 43 44 45 def call(self, inputs): 46 """ 47 Args: 48 inputs (Tensor): Input image of shape (B, 256, 256, 1) 49 50 Returns: 51 keypoints (Tensor): (B, num_keypoints, 2) normalized (x, y) keypoint coordinates. 52 jacobians (Tensor): (B, num_keypoints, 2, 2) local affine Jacobian matrices. 53 """ 54 batch_size = keras.ops.shape(inputs)[0] 55 56 key_feat = self.keypoint_branch(inputs) 57 jac_feat = self.jacobian_branch(inputs) 58 59 H_feat, W_feat = key_feat.shape[1], key_feat.shape[2] 60 grid = generate_grid((batch_size, H_feat, W_feat)) 61 62 key_feat_flat = ops.reshape(key_feat, (batch_size, -1, self.num_keypoints)) 63 weights = ops.softmax(key_feat_flat, axis=1) 64 weights = ops.reshape(weights, (batch_size, H_feat, W_feat, self.num_keypoints)) 65 66 grid_x = ops.expand_dims(grid[..., 0], axis=-1) 67 grid_y = ops.expand_dims(grid[..., 1], axis=-1) 68 kp_x = ops.sum(weights * grid_x, axis=[1, 2]) 69 kp_y = ops.sum(weights * grid_y, axis=[1, 2]) 70 keypoints = ops.stack([kp_x, kp_y], axis=-1) 71 72 jac_feat = ops.reshape(jac_feat, (batch_size, H_feat, W_feat, self.num_keypoints, 4)) 73 jac_feat = ops.transpose(jac_feat, [0, 4, 1, 2, 3]) 74 weights_exp = ops.expand_dims(weights, axis=1) 75 weighted_jac = jac_feat * weights_exp 76 weighted_jac = ops.transpose(weighted_jac, [0, 2, 3, 4, 1]) 77 weighted_jac = ops.reshape(weighted_jac, (batch_size, H_feat, W_feat, 4 * self.num_keypoints)) 78 jac_sum = ops.sum(weighted_jac, axis=[1, 2]) 79 jacobians = ops.reshape(jac_sum, (batch_size, self.num_keypoints, 2, 2)) 80 81 return keypoints, jacobians
11class KeypointDetector(keras.layers.Layer): 12 def __init__(self, num_keypoints=10, keypoint_branch=None, jacobian_branch=None, **kwargs): 13 """ 14 Layer for detecting keypoints and local Jacobian transformations. 15 16 Args: 17 num_keypoints (int): Number of keypoints to detect. 18 keypoint_branch (keras.Model or None): Optional custom model for keypoint detection. 19 If None, a default convolutional branch will be used. 20 jacobian_branch (keras.Model or None): Optional custom model for Jacobian prediction. 21 If None, a default convolutional branch will be used. 22 **kwargs: Additional keyword arguments passed to the parent class. 23 """ 24 super().__init__(**kwargs) 25 self.num_keypoints = num_keypoints 26 27 if keypoint_branch is None: 28 self.keypoint_branch = keras.Sequential([ 29 keras.layers.Conv2D(32, 3, strides=2, activation='relu', padding='same'), 30 keras.layers.Conv2D(32, 3, strides=2, activation='relu', padding='same'), 31 keras.layers.Conv2D(num_keypoints, 3, strides=1, activation=None, padding='same'), 32 ]) 33 else: 34 self.keypoint_branch = keypoint_branch 35 36 if jacobian_branch is None: 37 self.jacobian_branch = keras.Sequential([ 38 keras.layers.Conv2D(32, 3, strides=2, activation='relu', padding='same'), 39 keras.layers.Conv2D(32, 3, strides=2, activation='relu', padding='same'), 40 keras.layers.Conv2D(4 * num_keypoints, 3, strides=1, activation=None, padding='same'), 41 ]) 42 else: 43 self.jacobian_branch = jacobian_branch 44 45 46 def call(self, inputs): 47 """ 48 Args: 49 inputs (Tensor): Input image of shape (B, 256, 256, 1) 50 51 Returns: 52 keypoints (Tensor): (B, num_keypoints, 2) normalized (x, y) keypoint coordinates. 53 jacobians (Tensor): (B, num_keypoints, 2, 2) local affine Jacobian matrices. 54 """ 55 batch_size = keras.ops.shape(inputs)[0] 56 57 key_feat = self.keypoint_branch(inputs) 58 jac_feat = self.jacobian_branch(inputs) 59 60 H_feat, W_feat = key_feat.shape[1], key_feat.shape[2] 61 grid = generate_grid((batch_size, H_feat, W_feat)) 62 63 key_feat_flat = ops.reshape(key_feat, (batch_size, -1, self.num_keypoints)) 64 weights = ops.softmax(key_feat_flat, axis=1) 65 weights = ops.reshape(weights, (batch_size, H_feat, W_feat, self.num_keypoints)) 66 67 grid_x = ops.expand_dims(grid[..., 0], axis=-1) 68 grid_y = ops.expand_dims(grid[..., 1], axis=-1) 69 kp_x = ops.sum(weights * grid_x, axis=[1, 2]) 70 kp_y = ops.sum(weights * grid_y, axis=[1, 2]) 71 keypoints = ops.stack([kp_x, kp_y], axis=-1) 72 73 jac_feat = ops.reshape(jac_feat, (batch_size, H_feat, W_feat, self.num_keypoints, 4)) 74 jac_feat = ops.transpose(jac_feat, [0, 4, 1, 2, 3]) 75 weights_exp = ops.expand_dims(weights, axis=1) 76 weighted_jac = jac_feat * weights_exp 77 weighted_jac = ops.transpose(weighted_jac, [0, 2, 3, 4, 1]) 78 weighted_jac = ops.reshape(weighted_jac, (batch_size, H_feat, W_feat, 4 * self.num_keypoints)) 79 jac_sum = ops.sum(weighted_jac, axis=[1, 2]) 80 jacobians = ops.reshape(jac_sum, (batch_size, self.num_keypoints, 2, 2)) 81 82 return keypoints, jacobians
This is the class from which all layers inherit.
A layer is a callable object that takes as input one or more tensors and
that outputs one or more tensors. It involves computation, defined
in the call()
method, and a state (weight variables). State can be
created:
- in
__init__()
, for instance viaself.add_weight()
; - in the optional
build()
method, which is invoked by the first__call__()
to the layer, and supplies the shape(s) of the input(s), which may not have been known at initialization time.
Layers are recursively composable: If you assign a Layer instance as an
attribute of another Layer, the outer layer will start tracking the weights
created by the inner layer. Nested layers should be instantiated in the
__init__()
method or build()
method.
Users will just instantiate a layer and then treat it as a callable.
Args:
trainable: Boolean, whether the layer's variables should be trainable.
name: String name of the layer.
dtype: The dtype of the layer's computations and weights. Can also be a
keras.DTypePolicy
,
which allows the computation and
weight dtype to differ. Defaults to None
. None
means to use
keras.config.dtype_policy()
,
which is a float32
policy unless set to different value
(via keras.config.set_dtype_policy()
).
Attributes:
name: The name of the layer (string).
dtype: Dtype of the layer's weights. Alias of layer.variable_dtype
.
variable_dtype: Dtype of the layer's weights.
compute_dtype: The dtype of the layer's computations.
Layers automatically cast inputs to this dtype, which causes
the computations and output to also be in this dtype.
When mixed precision is used with a
keras.DTypePolicy
, this will be different
than variable_dtype
.
trainable_weights: List of variables to be included in backprop.
non_trainable_weights: List of variables that should not be
included in backprop.
weights: The concatenation of the lists trainable_weights and
non_trainable_weights (in this order).
trainable: Whether the layer should be trained (boolean), i.e.
whether its potentially-trainable weights should be returned
as part of layer.trainable_weights
.
input_spec: Optional (list of) InputSpec
object(s) specifying the
constraints on inputs that can be accepted by the layer.
We recommend that descendants of Layer
implement the following methods:
__init__()
: Defines custom layer attributes, and creates layer weights that do not depend on input shapes, usingadd_weight()
, or other state.build(self, input_shape)
: This method can be used to create weights that depend on the shape(s) of the input(s), usingadd_weight()
, or other state.__call__()
will automatically build the layer (if it has not been built yet) by callingbuild()
.call(self, *args, **kwargs)
: Called in__call__
after making surebuild()
has been called.call()
performs the logic of applying the layer to the input arguments. Two reserved keyword arguments you can optionally use incall()
are: 1.training
(boolean, whether the call is in inference mode or training mode). 2.mask
(boolean tensor encoding masked timesteps in the input, used e.g. in RNN layers). A typical signature for this method iscall(self, inputs)
, and user could optionally addtraining
andmask
if the layer need them.get_config(self)
: Returns a dictionary containing the configuration used to initialize this layer. If the keys differ from the arguments in__init__()
, then overridefrom_config(self)
as well. This method is used when saving the layer or a model that contains this layer.
Examples:
Here's a basic example: a layer with two variables, w
and b
,
that returns y = w . x + b
.
It shows how to implement build()
and call()
.
Variables set as attributes of a layer are tracked as weights
of the layers (in layer.weights
).
class SimpleDense(Layer):
def __init__(self, units=32):
super().__init__()
self.units = units
# Create the state of the layer (weights)
def build(self, input_shape):
self.kernel = self.add_weight(
shape=(input_shape[-1], self.units),
initializer="glorot_uniform",
trainable=True,
name="kernel",
)
self.bias = self.add_weight(
shape=(self.units,),
initializer="zeros",
trainable=True,
name="bias",
)
# Defines the computation
def call(self, inputs):
return ops.matmul(inputs, self.kernel) + self.bias
# Instantiates the layer.
linear_layer = SimpleDense(4)
# This will also call `build(input_shape)` and create the weights.
y = linear_layer(ops.ones((2, 2)))
assert len(linear_layer.weights) == 2
# These weights are trainable, so they're listed in `trainable_weights`:
assert len(linear_layer.trainable_weights) == 2
Besides trainable weights, updated via backpropagation during training,
layers can also have non-trainable weights. These weights are meant to
be updated manually during call()
. Here's a example layer that computes
the running sum of its inputs:
class ComputeSum(Layer):
def __init__(self, input_dim):
super(ComputeSum, self).__init__()
# Create a non-trainable weight.
self.total = self.add_weight(
shape=(),
initializer="zeros",
trainable=False,
name="total",
)
def call(self, inputs):
self.total.assign(self.total + ops.sum(inputs))
return self.total
my_sum = ComputeSum(2)
x = ops.ones((2, 2))
y = my_sum(x)
assert my_sum.weights == [my_sum.total]
assert my_sum.non_trainable_weights == [my_sum.total]
assert my_sum.trainable_weights == []
12 def __init__(self, num_keypoints=10, keypoint_branch=None, jacobian_branch=None, **kwargs): 13 """ 14 Layer for detecting keypoints and local Jacobian transformations. 15 16 Args: 17 num_keypoints (int): Number of keypoints to detect. 18 keypoint_branch (keras.Model or None): Optional custom model for keypoint detection. 19 If None, a default convolutional branch will be used. 20 jacobian_branch (keras.Model or None): Optional custom model for Jacobian prediction. 21 If None, a default convolutional branch will be used. 22 **kwargs: Additional keyword arguments passed to the parent class. 23 """ 24 super().__init__(**kwargs) 25 self.num_keypoints = num_keypoints 26 27 if keypoint_branch is None: 28 self.keypoint_branch = keras.Sequential([ 29 keras.layers.Conv2D(32, 3, strides=2, activation='relu', padding='same'), 30 keras.layers.Conv2D(32, 3, strides=2, activation='relu', padding='same'), 31 keras.layers.Conv2D(num_keypoints, 3, strides=1, activation=None, padding='same'), 32 ]) 33 else: 34 self.keypoint_branch = keypoint_branch 35 36 if jacobian_branch is None: 37 self.jacobian_branch = keras.Sequential([ 38 keras.layers.Conv2D(32, 3, strides=2, activation='relu', padding='same'), 39 keras.layers.Conv2D(32, 3, strides=2, activation='relu', padding='same'), 40 keras.layers.Conv2D(4 * num_keypoints, 3, strides=1, activation=None, padding='same'), 41 ]) 42 else: 43 self.jacobian_branch = jacobian_branch
Layer for detecting keypoints and local Jacobian transformations.
Args: num_keypoints (int): Number of keypoints to detect. keypoint_branch (keras.Model or None): Optional custom model for keypoint detection. If None, a default convolutional branch will be used. jacobian_branch (keras.Model or None): Optional custom model for Jacobian prediction. If None, a default convolutional branch will be used. **kwargs: Additional keyword arguments passed to the parent class.
46 def call(self, inputs): 47 """ 48 Args: 49 inputs (Tensor): Input image of shape (B, 256, 256, 1) 50 51 Returns: 52 keypoints (Tensor): (B, num_keypoints, 2) normalized (x, y) keypoint coordinates. 53 jacobians (Tensor): (B, num_keypoints, 2, 2) local affine Jacobian matrices. 54 """ 55 batch_size = keras.ops.shape(inputs)[0] 56 57 key_feat = self.keypoint_branch(inputs) 58 jac_feat = self.jacobian_branch(inputs) 59 60 H_feat, W_feat = key_feat.shape[1], key_feat.shape[2] 61 grid = generate_grid((batch_size, H_feat, W_feat)) 62 63 key_feat_flat = ops.reshape(key_feat, (batch_size, -1, self.num_keypoints)) 64 weights = ops.softmax(key_feat_flat, axis=1) 65 weights = ops.reshape(weights, (batch_size, H_feat, W_feat, self.num_keypoints)) 66 67 grid_x = ops.expand_dims(grid[..., 0], axis=-1) 68 grid_y = ops.expand_dims(grid[..., 1], axis=-1) 69 kp_x = ops.sum(weights * grid_x, axis=[1, 2]) 70 kp_y = ops.sum(weights * grid_y, axis=[1, 2]) 71 keypoints = ops.stack([kp_x, kp_y], axis=-1) 72 73 jac_feat = ops.reshape(jac_feat, (batch_size, H_feat, W_feat, self.num_keypoints, 4)) 74 jac_feat = ops.transpose(jac_feat, [0, 4, 1, 2, 3]) 75 weights_exp = ops.expand_dims(weights, axis=1) 76 weighted_jac = jac_feat * weights_exp 77 weighted_jac = ops.transpose(weighted_jac, [0, 2, 3, 4, 1]) 78 weighted_jac = ops.reshape(weighted_jac, (batch_size, H_feat, W_feat, 4 * self.num_keypoints)) 79 jac_sum = ops.sum(weighted_jac, axis=[1, 2]) 80 jacobians = ops.reshape(jac_sum, (batch_size, self.num_keypoints, 2, 2)) 81 82 return keypoints, jacobians
Args: inputs (Tensor): Input image of shape (B, 256, 256, 1)
Returns: keypoints (Tensor): (B, num_keypoints, 2) normalized (x, y) keypoint coordinates. jacobians (Tensor): (B, num_keypoints, 2, 2) local affine Jacobian matrices.