flax-u2net

• Flax/JAX implementation of U²-Net: Going Deeper with Nested U-Structure for Salient Object Detection.
• Original implementation found here
• Replaced batch normalization with layer normalization.
• Using Tensorflow Dataset API For data loading.
• DUTS dataset found here
• Attempt to generalize the RSU Block to avoid duplication of code for varying levels

Below is an example where the model detects the blini (a type of Russian pancake) instead of the cat.