arturml / pytorch-unet

Simple PyTorch U-Net for semantic segmentation of fish images.

Pytorch U-Net

This repository contains a simple PyTorch implementation of an U-Net for semantic segmentation of fish images, using this dataset by B. J. Boom, P. X. Huang and J. He, R. B. Fisher [1].

Here is a sample fish image and its ground truth mask:

The model is very simple and not super accurate, but the results are kinda cute:

The code for the U-Net is partially based on this Kaggle kernel.

[1] B. J. Boom, P. X. Huang, J. He, R. B. Fisher, "Supporting Ground-Truth annotation of image datasets using clustering", 21st Int. Conf. on Pattern Recognition (ICPR), 2012

pytorch-unet's Issues

Root directory of images

How to change root to images in FishDataset.py (33-35 lines) if I have:
dataset/images/Img1.png...Img2.png etc.
dataset/mask/ImgMask1.png...ImgMask2.png etc.

Do you know how to do it?