patrickisroman / u2net-tensorflow Goto Github PK

Tensorflow implementation of U^2-Net [...] for Salient Object Detection

Python 100.00%

u2net-tensorflow's Introduction

u2net-tensorflow

A tensorflow implementation of the U2-Net: Going Deeper with Nested U-Structure for Salient Object Detection using Keras & Functional API

Based on the PyTorch version by NathanUA, PDillis, vincentzhang, and chenyangh

Notebook

If you just want to play with the model, I've setup a Google Colab Notebook that lets you train the model on DUTS-TR, and it's fun to watch the model learn to mask an image of the Space Needle that it's never seen before while it trains. Training takes ~60 minutes to get to noticeable results, but you should train for several hours to use it for testing. Please let me know if you have any questions.

Setup

virtualenv venv
source venv/bin/activate
pip install tensorflow matplotlib opencv-python wget

Training

OPTIONAL: Download the DUTS-TR dataset and unzip into the data directory to load the training set:

wget http://saliencydetection.net/duts/download/DUTS-TR.zip
unzip DUTS-TR.zip -d data

If train.py does not detect the dataset is present when run, it will automatically try to download and setup the dataset before initiating training. If you have a custom dataset, you can update dataset_dir, image_dir, and mask_dir in config.py. Images will be rescaled

To begin training, simply run:

python train.py

Weights are automatically saved every save_interval iterations to weights/u2net.h5. These can be overwritten by passing the appropriate arguments. See python train.py -h for args.

Evaluation

Use eval.py to evaluate the model on images:

python eval.py --weights=weights/u2net.h5 --image=examples/skateboard.jpg

By default, the output images are saved in the out subdirectory

Custom Usage

The U2NET class can be used to instatiate a modular instance of the U2-Net network

from model.u2net import U2NET

u2net = U2NET()
out = u2net(inp)

u2net-tensorflow's People

Contributors

Stargazers

Watchers

u2net-tensorflow's Issues

OSError: Unable to open file (bad object header version number)

Thanks for such a nice work. I have trained the model using train.py file and save '.h5' model.Now I am testing the model using the 'eval.py' file.But I am getting the following error
Traceback (most recent call last):
File "eval.py", line 99, in
main()
File "eval.py", line 68, in main
model.load_weights(args.weights)
File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/training.py", line 2227, in load_weights
with h5py.File(filepath, 'r') as f:
File "/usr/local/lib/python3.6/dist-packages/h5py/_hl/files.py", line 408, in init
swmr=swmr)
File "/usr/local/lib/python3.6/dist-packages/h5py/_hl/files.py", line 173, in make_fid
fid = h5f.open(name, flags, fapl=fapl)
File "h5py/_objects.pyx", line 54, in h5py._objects.with_phil.wrapper
File "h5py/_objects.pyx", line 55, in h5py._objects.with_phil.wrapper
File "h5py/h5f.pyx", line 88, in h5py.h5f.open
OSError: Unable to open file (bad object header version number)

TypeError: 'NoneType' object is not subscriptable - when run train.py

Pretrained model for evaluation

has anyone got a trained model?Thanks

TypeError: 'NoneType' object is not callable metrics

How can I add metrics of evaluation in the train step?
i add metrics 'acc' and 'mae'
model = keras.Model(inputs=inputs, outputs=out, name='u2netmodel')
model.compile(optimizer=adam, loss=bce_loss, metrics=['acc', 'mae'])
when i run train: error:
TypeError Traceback (most recent call last)
in ()
4 try:
5 feed, out = load_training_batch(batch_size=batch_size)
----> 6 loss = model.train_on_batch(feed, out)
7
8 if e % weights_save_interval == 0:

2 frames
/usr/local/lib/python3.6/dist-packages/tensorflow/python/eager/def_function.py in _call(self, *args, **kwds)
853 # In this case we have created variables on the first call, so we run the
854 # defunned version which is guaranteed to never create variables.
--> 855 return self._stateless_fn(*args, **kwds) # pylint: disable=not-callable
856 elif self._stateful_fn is not None:
857 # Release the lock early so that multiple threads can perform the call

TypeError: 'NoneType' object is not callable
How can I fix this error?

Pyramid upsampling in the u-2net model

Hi
Thanks for sharing your code.
I saw that you use pyramid upsampling where you use the upsampling layer multiple times instead of single time for every side output tensor. Wanted to ask why did you choose to implement like so?
Thanks
Adva

When training the model, are the dimensions uniformly equal to 320 for width and 320 for height?

Cant save entire model

Hi i want to save entire keras model so that it can be converted to tflite. but I am getting "Notimplemted error". Any work around for that?
I am using model.save(model_path) to save keras model

Confusion in Weights h5 or ht5

@patrickisroman Thank you for nice work. I am using your Google Colab notebook for training. I am confused in weights saving format ".h5" and ".ht5" weights.You are saving weights as ".ht5" in colab but in "eval.py" you are using the .h5 foarmat for inferences. Can you guide me about the weights format .thanks