SemanticModel

Deep learning framework for semantic segmentation using PyTorch.

Install

pip install -r requirements.txt
python setup.py install

Usage

from SemanticModel.model_core import SegmentationModel
from SemanticModel.prediction import PredictionPipeline

# Train
model = SegmentationModel(
    classes=['background', 'object'],
    architecture='unet',
    encoder='timm-regnety_120' 
)

trainer = ModelTrainer(
    model_config=model,
    root_dir='path/to/dataset',
    epochs=40
)
model, metrics = trainer.train()

# Predict 
predictor = PredictionPipeline(model)
predictor.predict_single_image('image.jpg')
predictor.predict_directory('image_dir/')
predictor.predict_raster('raster.tif')

# Load pretrained
model = SegmentationModel(
    classes=['background', 'object'],
    weights='path/to/best_model.pth'
)

Data Structure

dataset/
├── train/
│   ├── Images/
│   └── Masks/
└── val/
    ├── Images/ 
    └── Masks/