This repo comprises of the training, inference and model conversion (pytorch model to onnx model) code for UNet masking model using Xircuits.
This repository uses the leeds butterfly dataset as its example. You may download it here and place it in your working directory, or let the PrepareButterflyDataset component do it for you.
pip install -r requirements.txt
Description
- This repo comprises of the training, inference and model conversion (pytorch model to onnx model) code for UNet masking model. It contains 9 components:
- ConvertTorchModelToOnnx
- CreateUnetModel
- ImageTrainTestSplit
- PrepareUnetDataLoader
- TrainUnet
- UNetModel
- UnetPredict
- ReadMaskDataset
- PrepareButterflyDataset
With 3 workflows:
- PyTorchUnetInferenceSample.xircuits - To allow inferencing of trained Unet model in either pytorch or onnx format.
- PyTorchUnetTrainSample.xircuits - To allow training of binary class masking with Unet model.
- PyTorchToOnnxSample.xircuits - To allow conversion of torch model to onnx model.
To Test
- To perform the training in PyTorchUnetTrainSample.xircuits. Please download any mask dataset that contains the images and segmentation images.
- The dataset that I used would be Leeds butterfly dataset and ensure the tree follows as below:
- The trained model should be saved in the examples folder by default but you are allowed to modify the save path.
- To allow it to work, you must have the same image name with the image and mask image. The prepare butterfly dataset component will do the setup for you.
-
To perform the inference code in PyTorchUnetInferenceSample.xircuits. Please put the path that model has be saved as model_path and inference image as image_path.
-
To perform the model conversion code in PytorchToOnnxSample.xircuits. Please put the pth format model as input_model_path and onnx format path as output_model_path.
Reference: Citing Leeds buttefly dataset Josiah Wang, Katja Markert, and Mark Everingham Learning Models for Object Recognition from Natural Language Descriptions In Proceedings of the 20th British Machine Vision Conference (BMVC2009)