Add BASNet to keras hub

[laxmareddyp]

Refer this issue for more details

Please refer working demo notebook

[divyashreepathihalli]

Thanks for the PR Laxma! it looks great! can you add a demo colab to verify outputs? Thanks!

[mattdangerw]

Started a review. But I think this might need some bigger considerations.

Some questions:

• What are the weight sources we are using here?
• What are the outputs? Looks like there is a list of outputs instead of a normal classification style output.
• How does training this model look?

In general, I think there there will be a lot of core infrastructure that will break if the arch does not have a backbone class. DeepLabV3 is probably much closer to what we will need to go for.

But I think I need answers to the above questions to better suggest a design.

[laxmareddyp]

Some questions:

• What are the weight sources we are using here?

The weights conversion script uses the weights from https://keras.io/examples/vision/basnet_segmentation/ . There's also a Keras BASNet model available on Kaggle, but it failed to load for me so far, and I don't know if where it comes from, if it's better or if it's the same model anyway.

• What are the outputs? Looks like there is a list of outputs instead of a normal classification style output?

Yes, the direct model output is a list for refinement and prediction modules. The different stages of the model are trained with the same segmentation mask data. For inference, access to different stages of prediction modules would not be required for the user, but it might be interesting and doesn't hurt.

• How does training this model look?

Training is simply performed on segmentation map data. For reference check this keras-cv notebook published in keras-io examples and issue

[mattdangerw]

Maybe for now just directly subclass resnet without changes as BasnetBackbone? That way we can satisfy all the infra that will expect every new arch to have a backbone class.

The broader question I have is how to make our ImageSegmenter tasks interchangeable in terms of data in and prediction out. The goal of the task API is to allow users to swap different arches while keep the overall training/inference code the same. But it seems like there's a number of different ways to pose the segmentation problem. Not sure what to do here, but we can at least add a backbone class for now.

[laxmareddyp]

Maybe for now just directly subclass resnet without changes as BasnetBackbone? That way we can satisfy all the infra that will expect every new arch to have a backbone class.

The broader question I have is how to make our ImageSegmenter tasks interchangeable in terms of data in and prediction out. The goal of the task API is to allow users to swap different arches while keep the overall training/inference code the same. But it seems like there's a number of different ways to pose the segmentation problem. Not sure what to do here, but we can at least add a backbone class for now.

I have created a BASNetBackbone and moved the architecture implementation there (similar to how it is done in DeepLab). BASNetImageSegmenter now outputs a single segmentation map, therefore should be compatible to your described interface. I've implemented a compute_loss() to train BASNet's prediction modules as well.

[mattdangerw]

Nice! Will take a look. Test failures look to be legit and basnet related, please take a look!

[mattdangerw]

Let's try this out! I wonder if we will need to make more tweaks, having the num classes on the backbone itself is kinda funky, but do get why we are doing it for this architecture.

Left a number of style comments.

[mattdangerw]

@laxmareddyp check the test failures! ~6 lines are too long and need fixing

[laxmareddyp]

@laxmareddyp check the test failures! ~6 lines are too long and need fixing

@mattdangerw I am modifying the code related to doc strings, other minor suggestions you mentioned and will push the changes and hope it will not show up this time. Thanks

[laxmareddyp]

Let's try this out! I wonder if we will need to make more tweaks, having the num classes on the backbone itself is kinda funky, but do get why we are doing it for this architecture.

num_classes: Unlike DeepLab, in BASNet this parameter is quite involved in the architecture. This is because BASNet has multiple model outputs, Therefore, we cannot construct the task by just adding a conv layer with num_classes outputs to the backbone. A possibility to achieve this goal would be splitting the backbone in several parts, which are used individually in the task, but the effort for this would (in my opinion) outweigh the benefits by far.

[mattdangerw]

Looks good overall. I'll drop the conversion script as I merge. Let's do that separately, need some work

[mattdangerw]

lgtm!