YOLO v3 and MiDaS from a single Resnext101 backbone.
By: Saravana Alagar
Tags: Deep Learning, Depth Estimation, Yolov3, Midas, Object Detection

Approach

MiDaS was a network which is really trained well. The weights can be used as it is. Connect 3 YOLO layers with the MiDaS network’s Feature Extractor. Add some convolution layers before YOLO layers in order to learn the Bounding Boxes better. Freeze the MiDas branch and train only for the YOLO branch, save weights. Introduce loss function for MiDaS network and define the total loss. Load the model with all pre-trained weights and train all layers with a very less learning rate.

Resnext101 with Imagenet weights is a very capable network for the Object Detection problem, moreover, the Darknet backbone used in YOLOv3 is in fact a variant of Resnet. Hence training the backbone from scratch is not a good idea especially when two branches are using it. Hence we add the capability to learn Bounding Boxes after the Backbone and do transfer learning.

Stitching YOLO with MiDaS

Now let’s train this network, as you can see during this exercise the MiDaS Decoder is frozen.

Training only the YOLO brach for 50 epochs

After just 50 epochs we are seeing an amazing result! Remember the original YOLO v3 was trained for 300 epochs. Now we save the model weights and focus on the loss function for the total network including MiDaS.

Loss Function

We do not have a detailed explanation on training a MiDaS network from scratch or otherwise it is difficult to train such a network with limited resources at our disposal hence we use the pre-trained MiDaS and generate depth maps for our images.

original images

Depth maps inferred from MiDaS

We use these depth maps as the ground truth for our Network. Now while training the network from both the branches the gradients from each of the braches could seep into the backbone weights hence might spoil them completely. First, let us make sure we define a loss function to the MiDaS Network and train it separately to see if the loss reduces(the loss will reduce very very minimal since we are already using a well-trained network).

We can use image quality assessment as a metric for loss. We can either use

SSIM (Structural Similarity Index) MSE (Pixelwise)

SSIM (proposed to be closer to human perception) measures the structural similarity between images in terms of luminance, contrast, and structures. SSIM will be 1 between like images and 0 between random images. But in our observation, we see an SSIM value of 0.00711 for the below two depth images

SSIM 0.00711

According to our problem, we need a metric that could tell us these two images are somewhat similar to each other. We will choose Pixel-wise MSE as a loss function because of its simplicity. With this, we can train our MiDaS branch alone freezing the YOLO branch.

Loss reducing for the depth map prediction

Now let us define the total loss for the network.

After training with the above-mentioned loss function with a very less learning rate(in order to not nudge the weights too much off the range as individually the branches performed quite good) we get the results as below.

This looks pretty good, now we can use this network to deploy in a navigation system. Thanks for Reading!

References: