75 lines
3.0 KiB
Python
75 lines
3.0 KiB
Python
"""
|
|
Trainer to learn depth information on unlabeled data (raw images/videos)
|
|
|
|
Allows pluggable depth networks for differing performance (including fast-depth)
|
|
"""
|
|
|
|
import tensorflow.keras as keras
|
|
import warp
|
|
import unsupervised.loss as loss
|
|
|
|
|
|
class UnsupervisedPoseDepthLearner(keras.Model):
|
|
"""
|
|
Keras model to learn simultaneous depth + pose from image/video sequences.
|
|
|
|
To train this, the datasource should yield 3 frames and camera intrinsics.
|
|
Optionally velocity + timestamp per frame to train to real scale
|
|
"""
|
|
|
|
def __init__(self, depth_model, pose_model, num_scales=3, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
self.depth_model = depth_model
|
|
self.pose_model = pose_model
|
|
self.num_scales = num_scales
|
|
|
|
|
|
def train_step(self, data):
|
|
"""
|
|
|
|
:param data: Format: {frames: Mat[3], intrinsics: Tensor}
|
|
"""
|
|
# Pass through depth for target image
|
|
# TODO: Convert frame to tensor (or do this in the dataloader)
|
|
# TODO: Ensure the depth output includes enough outputs for each scale
|
|
depth = self.depth_model(data.frames[1])
|
|
|
|
# Pass through depth -> pose for both source images
|
|
# TODO: Concat these poses using tf.concat
|
|
pose1 = self.pose_model(data.frames[1], data.frames[0])
|
|
pose2 = self.pose_model(data.frames[1], data.frames[2])
|
|
|
|
shape = depth[0].shape
|
|
|
|
# TODO: Pull coords out of train step into initialiser, then it only needs to be created once.
|
|
# Ideally the size/batch size will still be calculated automatically
|
|
coords = warp.image_coordinate(shape[0], shape[1], shape[2])
|
|
|
|
scale_losses = []
|
|
# For each scale, do the projective inverse warp step and calculate losses
|
|
for i in range(self.num_scales):
|
|
# TODO: Could simplify this by stacking the source images (see sfmlearner)
|
|
# It isn't too much of an issue right now since we're only using 2 images like in monodepth
|
|
# For each depth output (scale), do the projective inverse warp on each input image and calculate the losses
|
|
# Only take the min loss between the two warped images (from monodepth2)
|
|
warp1 = warp.projective_inverse_warp(data.frames[0], depth[i], pose1, data.intrinsics, coords)
|
|
warp2 = warp.projective_inverse_warp(data.frames[2], depth[i], pose1, data.intrinsics, coords)
|
|
|
|
# Per pixel loss is just the difference in pixel intensities?
|
|
# Something like l1 plus ssim
|
|
loss1 = loss.make_combined_ssim_l1_loss(data.frames[1], warp1)
|
|
loss2 = loss.make_combined_ssim_l1_loss(data.frames[1], warp2)
|
|
|
|
# Take the min from these? Or min after auto-masking? I think after auto masking
|
|
|
|
|
|
# Also do the auto masking from monodepth2 (compare pixel difference between warped with difference
|
|
# in source, if source is more different then ignore the pixel).
|
|
pass
|
|
|
|
# Collect losses, average them out
|
|
|
|
# Calculate smooth losses
|
|
|
|
|
|
pass |