search

davidtvs / PyTorch-ENet

PyTorch implementation of ENet
MIT License
389 stars 129 forks source link

Custom dataset #6

Closed CalinTimbusLucian closed 6 years ago

CalinTimbusLucian commented 6 years ago

I was wondering what it the solution in case we want to train it on our own dataset? Thank you in advance.

davidtvs commented 6 years ago

Try the following:

  1. The first step would be to create the dataloader. You can find examples in data/camvid.py and data/cityscapes.py. An important detail is that two PIL images, the sample image and the target image, must be returned by the __getitem__ function.
  2. Add your dataloader to data/__init__.py.
  3. In args.py you should add a new option to the --dataset command-line argument.
  4. In main.py you handle that option by importing the dataloader.
  5. Run the commands found in the README with your dataset --dataset YourDatasetName
CalinTimbusLucian commented 6 years ago

Thank you for the quick response! Really appreciate it!