Implementation of various semantic segmentation models in tensorflow & keras including popular datasets
Project description
TF Semantic Segmentation
Features
-
Datasets
- Ade20k
- Camvid
- Cityscapes
- MappingChallenge
- MotsChallenge
- Coco
- PascalVoc2012
- Taco
- Shapes (randomly creating triangles, rectangles and circles)
- Toy (Overlaying TinyImageNet with MNIST)
- ISIC2018
- CVC-ClinicDB
-
Distributed Training on Multiple GPUs
-
Hyper Parameter Optimization using WandB
-
WandB Integration
-
Easily create TFRecord from Directory
-
Tensorboard visualizations
-
Models:
- Unet
- Erfnet
- MultiResUnet
-
Losses:
- Catagorical Crossentropy
- Binary Crossentropy
- Crossentropy + SSIM
- Dice
- Crossentropy + Dice
- Tversky
- Focal
- Focal + Tversky
-
Activations:
- mish
- swish
- relu6
-
Optimizers:
- Ranger
- RAdam
-
Normalization
- Instance Norm
- Batch Norm
-
On the fly Augmentations
- flip left/right
- flip up/down
- rot 180
- color
Requirements
sudo apt-get install libsm6 libxext6 libxrender-dev libyaml-dev libpython3-dev
Tensorflow (2.x) & Tensorflow Addons (optional)
pip install tensorflow-gpu==2.1.0 --upgrade
pip install tensorflow-addons==0.7.0 --upgrade
Training
Hint: To see train/test/val images you have to start tensorboard like this
tensorboard --logdir=logs/ --reload_multifile=true
On inbuild datasets (generator)
python -m tf_semantic_segmentation.bin.train -ds 'tacobinary' -bs 8 -e 100 \
-logdir 'logs/taco-binary-test' -o 'adam' -lr 5e-3 --size 256,256 \
-l 'binary_crossentropy' -fa 'sigmoid' \
--train_on_generator --gpus='0' \
--tensorboard_train_images --tensorboard_val_images
Using a fixed record path
python -m tf_semantic_segmentation.bin.train --record_dir=records/cityscapes-512x256-rgb/ \
-bs 4 -e 100 -logdir 'logs/cityscapes-bs8-e100-512x256' -o 'adam' -lr 1e-4 -l 'categorical_crossentropy' \
-fa 'softmax' -bufsize 50 --metrics='iou_score,f1_score' -m 'erfnet' --gpus='0' -a 'mish' \
--tensorboard_train_images --tensorboard_val_images
Multi GPU training
python -m tf_semantic_segmentation.bin.train --record_dir=records/cityscapes-512x256-rgb/ \
-bs 4 -e 100 -logdir 'logs/cityscapes-bs8-e100-512x256' -o 'adam' -lr 1e-4 -l 'categorical_crossentropy' \
-fa 'softmax' -bufsize 50 --metrics='iou_score,f1_score' -m 'erfnet' --gpus='0,1,2,3' -a 'mish'
Using Code
from tf_semantic_segmentation.bin.train import train_test_model, get_args
# get the default args
args = get_args({})
# change some parameters
# !rm -r logs/
args.model = 'erfnet'
# args['color_mode'] = 0
args.batch_size = 8
args.size = [128, 128] # resize input dataset to this size
args.epochs = 10
args.learning_rate = 1e-4
args.optimizer = 'adam' # ['adam', 'radam', 'ranger']
args.loss = 'dice'
args.logdir = 'logs'
args.record_dir = "datasets/shapes/records"
args.final_activation = 'softmax'
# train and test
results, model = train_test_model(args)
Models
- Erfnet
- Unet
from tf_semantic_segmentation import models
# print all available models
print(list(modes.models_by_name.keys()))
# returns a model without the final activation function
# because the activation function depends on the loss function
model = models.get_model_by_name('erfnet', {"input_shape": (128, 128, 3), "num_classes": 5})
# call models directly
model = models.erfnet(input_shape=(128, 128), num_classes=5)
Use your own dataset
- Accepted file types are: jpg(jpeg) and png
If you already have a train/test/val split then use the following data structure:
dataset/
labels.txt
test/
images/
masks/
train/
images/
masks/
val/
images/
masks/
or use
dataset/
labels.txt
images/
masks/
The labels.txt should contain a list of labels separated by newline [/n]. For instance it looks like this:
background
car
pedestrian
- To create a tfrecord using the original image size and color use the script like this:
INPUT_DIR = ...
tf-semantic-segmentation-tfrecord-writer -dir $INPUT_DIR -r $INPUT_DIR/records
There are the following addition arguments:
- -s [--size] '$width,$height' (f.e. "512,512")
- -rm [--resize_method] ('resize', 'resize_with_pad', 'resize_with_crop_or_pad)
- cm [--color_mode] (0=RGB, 1=GRAY, 2=NONE (default))
Datasets
from tf_semantic_sementation.datasets import get_dataset by name, datasets_by_name, DataType, get_cache_dir
# print availiable dataset names
print(list(datasets_by_name.keys()))
# get the binary (waste or not) dataset
data_dir = '/hdd/data/'
name = 'tacobinary'
cache_dir = get_cache_dir(data_dir, name.lower())
ds = get_dataset_by_name(name, cache_dir)
# print labels and classes
print(ds.labels)
print(ds.num_classes)
# print number of training examples
print(ds.num_examples(DataType.TRAIN))
# or simply print the summary
ds.summary()
TFRecords
This library simplicifies the process of creating a tfrecord dataset for faster training.
Write tfrecords:
from tf_semantic_segmentation.datasets import TFWriter
ds = ...
writer = TFWriter(record_dir)
writer.write(ds)
writer.validate(ds)
or use simple with this script (will be save with size 128 x 128 (width x height)):
tf-semantic-segmentation-tfrecord-writer -d 'toy' -c /hdd/datasets/ -s '128,128'
Docker
docker build -t tf_semantic_segmentation -f docker/Dockerfile ./
Prediction
pip install matplotlib
Using Code
from tensorflow.keras.models import load_model
import numpy as np
from tf_semantic_segmentation.processing import dataset
from tf_semantic_segmentation.visualizations import show, masks
model = load_model('logs/model-best.h5', compile=False)
# model parameters
size = tuple(model.input.shape[1:3])
depth = model.input.shape[-1]
color_mode = dataset.ColorMode.GRAY if depth == 1 else dataset.ColorMode.RGB
# define an image
image = np.zeros((256, 256, 3), np.uint8)
# preprocessing
image = image.astype(np.float32) / 255.
image, _ = dataset.resize_and_change_color(image, None, size, color_mode, resize_method='resize')
image_batch = np.expand_dims(image, axis=0)
# predict (returns probabilities)
p = model.predict(image_batch)
# draw segmentation map
num_classes = p.shape[-1] if p.shape[-1] > 1 else 2
predictions_rgb = masks.get_colored_segmentation_mask(p, num_classes, images=image_batch, binary_threshold=0.5)
# show images using matplotlib
show.show_images([predictions_rgb[0], image_batch[0]])
Using scripts
- On image
python -m tf_semantic_segmentation.evaluation.predict -m model-best.h5 -i image.png
- On TFRecord (data type 'val' is default)
python -m tf_semantic_segmentation.evaluation.predict -m model-best.h5 -r records/camvid/
- On TFRecord (with export to directory)
python -m tf_semantic_segmentation.evaluation.predict -m model-best.h5 -r records/cubbinary/ -o out/ -rm 'resize_with_pad'
- On Video
python -m tf_semantic_segmentation.evaluation.predict -m model-best.h5 -v video.mp4
- On Video (with export to out/p-video.mp4)
python -m tf_semantic_segmentation.evaluation.predict -m model-best.h5 -v video.mp4 -o out/
Project details
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
Hashes for tf_semantic_segmentation-0.2.2.tar.gz
Algorithm | Hash digest | |
---|---|---|
SHA256 | 2b4ad6c529d2945e291be8416cf3c0caeb3e31b2d469f1317b31df11c65fdd73 |
|
MD5 | 5ffcdc5dc5d4c44a7d13a5b3fdf4a46e |
|
BLAKE2b-256 | bfbe33403006d098b4e41e9a78c222577616abb92aa7e441ca8742f92895f5c7 |
Hashes for tf_semantic_segmentation-0.2.2-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | 835334c67b6152c45f20d08e3312f217cc7b31f49f7ce036aefd5f53951052c9 |
|
MD5 | ee1d07af5aa7bb4ff7b89a4aeaf6773e |
|
BLAKE2b-256 | d25aace3a95d4891f8765dc1aa655204842744d6f22e65adff8d6003caf4ab1e |