dataset/
:train_dataset.py
: The training and valid dataset script.test_dataset.py
: The test dataset script.transforms.py
: Useful transforms for dataset processing.
model/
:UNet.py
: The 3D U-Net model.
utils/
: The useful scripts.result/
:best_model.pth
: The model used for predict.test_predict.zip
: The test dataset prediction results.val_predict.zip
: The val dataset prediction results.
train.py
: The training script.predict.py
: The script to predict the label.config.py
: The script to get the parameters required for training.plot.py
: The script to plot the needed curve.
python train.py --lr 0.001 --gpu_id=[0,1,2,3] --batch_size 8 --train_image_dir 'datasets/ribfrac-train-images/' --train_label_dir 'datasets/ribfrac-train-labels/' --val_image_dir 'datasets/ribfrac-val-images/' --val_label_dir 'datasets/ribfrac-val-labels/'
lr:
learning rate (default: 0.001)gpu_id:
gpu id list used for training.batch_size:
batch size of trainsettrain_image_dir:
The folder where the training images are saved.train_label_dir:
The folder where the training labels are saved.val_image_dir:
The folder where the val images are saved.val_label_dir:
The folder where the val labels are saved.
For more parameter information, please check the config.py
script.
python predict.py --gpu 0 --model_path ./result/best_model.pth --pred_dir ./predict --test_data_path ./datasets/ribfrac-test-images/
gpu:
gpu id used for predict.model_path:
The model path used for predict.pred_dir:
The folder to store the predict result.test_data_path:
The folder where the test images are saved.
For more parameter information, please check the config.py
script.