3D U-Net Training Pipeline for Brain Cell Segmentation

This repository contains a modular and reproducible MATLAB pipeline for training a 3D U-Net model on volumetric TIFF images. It was developed for brain-wide cell-type segmentation using tissue-cleared images and pixel-level annotations.

---

Summary

This project enables patch-based training of a 3D U-Net using custom-labeled datasets. It supports Dice loss optimization and includes validation patch sampling.

• Input format: 3D TIFF stacks for raw images and pixel-level labels
• Output: Trained 3D U-Net model for voxel-wise segmentation
• Training method: Patch extraction and mini-batch training with Dice loss

---

Repository Structure

File	Description
prepare_data.m	Loads raw and label images, constructs training/validation patch datasets
define_model.m	Defines the 3D U-Net model architecture with encoder depth 3
train_model.m	Trains the model using Adam optimizer and Dice loss
run_all.m	Wrapper function that runs the complete pipeline

---

Requirements

• MATLAB R2021a or later
• Image Processing Toolbox
• Deep Learning Toolbox
• Access to GPU recommended for training

---

Input Directory Structure

/your_dataset/
├── training_raw/           # 3D TIFF files with raw input volumes
├── training_labels/        # Corresponding label TIFFs with pixel values [0, 127, 255]
├── validation_raw/         # Validation images
├── validation_labels/      # Validation labels

Ensure that file names in each pair of folders match one-to-one.

---

How to Run

1. In MATLAB:

run_all();

2. From Terminal (Linux/Mac):

matlab -batch "run_all"

---

Notes on Label Encoding

The labels must use the following grayscale pixel values:

Label	Pixel Value
background	0
outercell	127
cell	255

These are mapped in prepare_data.m and used by pixelLabelDatastore.