🧠 U-Net for Image Segmentation

This project provides a complete implementation of the U-Net architecture for image segmentation tasks. It trains the model using a dataset of input images and corresponding segmentation masks, which may be binary (e.g., foreground/background) or multi-class (e.g., different object categories).

U-Net is a convolutional neural network originally designed for biomedical image segmentation, but it is widely used across various domains for its accuracy and efficiency. This implementation includes configurable input preprocessing, model training, evaluation, augmentation, and deployment support.

---

📦 Dependencies

Install all required packages:

pip install -r requirements.txt

# 🧪 Data Augmentation

To increase dataset diversity and prevent overfitting, this project uses the Albumentations library for data augmentation. The augmentations are applied consistently to both input images and their corresponding masks.

▶️ How to Run Augmentation
Edit the augmentation pipeline inside augment.py as needed.

Then run:

python augment.py

# 🏋️ Training the Model
After preparing your dataset (original + augmented), run the training script:

python train.py

You can configure parameters such as:
Image size
Learning rate
Number of epochs
Batch size

All configurations are available inside train.py.

# 🧮 INT8 Quantization (Post-Training Optimization)

To deploy the trained U-Net model efficiently on edge devices (like Raspberry Pi, Coral Edge TPU, or mobile devices), INT8 quantization is supported using TensorFlow Lite.

Benefits:

Reduced model size
Faster inference
Minimal loss in accuracy

🔧 Run Quantization
Once the model is trained, run:
python int8_quantization.py

This script will generate a TFLite model in INT8 format.

# 🤖 Edge TPU Compilation (For Google Coral)
To deploy on Google Coral, the quantized TFLite model must be compiled using the Edge TPU Compiler.

🔍 Install Edge TPU Compiler (Linux)

echo "deb https://packages.cloud.google.com/apt coral-edgetpu-stable main" | sudo tee /etc/apt/sources.list.d/coral-edgetpu.list
curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo apt-key add -
sudo apt update
sudo apt install edgetpu-compiler

✅ Compile the Model
edgetpu_compiler unet_model_int8.tflite

If successful, output will be:
Model compiled successfully.
Output: model_int8_edgetpu.tflite

# 🧪 Run Inference on Edge TPU
After compilation, you can run inference using:

python inference.py 
Customize the script to suit your input/output processing needs.