Tomato Segmentation with U-Net

A deep learning project for semantic segmentation of tomato images using a U-Net architecture. This implementation is designed to run on SLURM-based computing clusters and uses PyTorch for model training and inference.

Project Overview

This project implements a U-Net model for precise tomato segmentation in images. The model is trained on the Laboro Tomato dataset and includes features such as:

• Custom dataset handling with automatic image orientation correction
• Data augmentation for improved model robustness
• SLURM-based training and testing scripts for HPC environments
• Comprehensive evaluation metrics including IoU (Intersection over Union)

Repository Setup

Clone the repository:

git clone https://github.com/jan3zk/tomato_segmentation.git
cd tomato_segmentation/

Dataset

Download the dataset using gdown and unzip it so the laboro_tomato directory appears in the project root:

gdown --id 1K5Zo47YIddzO3OnLgOJJbQig8TC_uSjb
unzip laboro_tomato.zip

If gdown is not installed you can add it with:

pip install --user gdown
# or using conda
conda install -c conda-forge gdown

The dataset follows the structure described below:

laboro_tomato/
├── train/
│   ├── images/tomato/*.png    # Training images
│   ├── masks/tomato/*.png     # Training masks
│   └── gt/tomato/*.png        # Ground truth masks
├── test/
│   ├── images/tomato/*.png    # Testing images
│   ├── masks/tomato/*.png     # Testing masks
│   └── gt/tomato/*.png        # Ground truth masks
└── annotations.json            # Bounding box and ripeness labels

By default train.py and test.py look for the laboro_tomato directory in the project root. Ensure the dataset was unzipped inside this repository or adjust the scripts if you keep it elsewhere.

Running in batch mode on HPC

Submit the training job:

sbatch train.sh

The training script:

• Loads the dataset from the laboro_tomato directory.
• Trains the U-Net model.
• Saves the best model checkpoint in the output_<timestamp>/checkpoints directory.

Submit the testing job with the model checkpoint path:

sbatch test.sh output_<timestamp>/checkpoints/best_model.pt

The testing script:

• Loads the trained model.
• Evaluates on the test dataset.
• Saves predictions and evaluation results in a test_results_<timestamp> directory.

Running in interactive mode on HPC

Start an interactive session on the cluster and load the required modules:

module purge
module load CUDA/11.7.0
module load PyTorch/1.12.0-foss-2022a-CUDA-11.7.0
module load Python/3.10.4-GCCcore-11.3.0
module load tqdm/4.64.1-GCCcore-12.2.0
module load matplotlib/3.5.2-foss-2022a

# Install any missing Python packages
pip install --user albumentations

From the repository directory run the training or testing script directly:

python train.py        # or python test.py --model_path <checkpoint>

Model Architecture

The U-Net model includes:

• Encoder: Extracts image features using convolutional layers.
• Decoder: Reconstructs the segmentation mask with skip connections.